Skip to content

Commit

Permalink
Refactor lib functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@lhvy
lhvy committed Nov 4, 2023
1 parent dfc00c5 commit b51e646
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 106 deletions.
109 changes: 9 additions & 100 deletions src/lib.rs → src/lsb.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use image::ImageBuffer;
use rand::seq::SliceRandom;
use rand::SeedableRng;

pub fn lsb_encode(message: String, mut image: image::RgbImage, n_bits: u8, key: Option<u64>) {
pub(crate) fn encode(message: String, mut image: image::RgbImage, n_bits: u8, key: Option<u64>) {
// Convert the message to a vector of bytes
let mut message = message.as_bytes().to_vec();
// Add another byte for optimisation purposes (required or data will be lost 💀)
Expand Down Expand Up @@ -33,14 +33,14 @@ pub fn lsb_encode(message: String, mut image: image::RgbImage, n_bits: u8, key:
shuffled_carrier.push(carrier[*i]);
}

lsb_raw_encode(&message, &mut shuffled_carrier, n_bits as usize);
raw_encode(&message, &mut shuffled_carrier, n_bits as usize);

for (i, shuffled_carrier) in index_map.into_iter().zip(shuffled_carrier) {
carrier[i] = shuffled_carrier;
}
}
None => {
lsb_raw_encode(&message, carrier, n_bits as usize);
raw_encode(&message, carrier, n_bits as usize);
}
}

Expand All @@ -49,7 +49,7 @@ pub fn lsb_encode(message: String, mut image: image::RgbImage, n_bits: u8, key:
println!("Saved image to ./out.png");
}

pub fn lsb_raw_encode(payload: &[u8], carrier: &mut [u8], n_bits: usize) {
fn raw_encode(payload: &[u8], carrier: &mut [u8], n_bits: usize) {
let len = u32::try_from(payload.len() - 1).unwrap();

// Write the length of the message to the LSB of the first 32 bytes
Expand Down Expand Up @@ -84,7 +84,7 @@ pub fn lsb_raw_encode(payload: &[u8], carrier: &mut [u8], n_bits: usize) {
}
}

pub fn lsb_decode(image: image::RgbImage, n_bits: u8, key: Option<u64>) {
pub(crate) fn decode(image: image::RgbImage, n_bits: u8, key: Option<u64>) {
// Decode the message from the image
let message = match key {
Some(key) => {
Expand All @@ -95,17 +95,17 @@ pub fn lsb_decode(image: image::RgbImage, n_bits: u8, key: Option<u64>) {

carrier.shuffle(&mut rng);

lsb_raw_decode(&carrier, n_bits as usize)
raw_decode(&carrier, n_bits as usize)
}
None => lsb_raw_decode(image.as_flat_samples().as_slice(), n_bits as usize),
None => raw_decode(image.as_flat_samples().as_slice(), n_bits as usize),
};

// Convert the message to a string
let message = String::from_utf8(message).unwrap();
println!("Message: {}", message);
}

pub fn lsb_raw_decode(carrier: &[u8], n_bits: usize) -> Vec<u8> {
fn raw_decode(carrier: &[u8], n_bits: usize) -> Vec<u8> {
// Read the length of the message from the LSB of the first 32 bytes
let mut len = 0;
for (i, byte) in carrier[..32].iter().enumerate() {
Expand Down Expand Up @@ -150,7 +150,7 @@ pub fn lsb_raw_decode(carrier: &[u8], n_bits: usize) -> Vec<u8> {
payload
}

pub fn extract_bit_planes(image: image::RgbImage) {
pub(crate) fn extract_bit_planes(image: image::RgbImage) {
let mut planes = vec![ImageBuffer::new(image.width(), image.height()); 24].into_boxed_slice();
let mut plane_pixels = planes
.iter_mut()
Expand Down Expand Up @@ -186,94 +186,3 @@ pub fn extract_bit_planes(image: image::RgbImage) {
}
println!("Saved images to ./output");
}

pub fn bytes_to_wav(bytes: &[u8]) {
let spec = hound::WavSpec {
channels: 1,
sample_rate: 44100,
bits_per_sample: 16,
sample_format: hound::SampleFormat::Int,
};

let mut writer = hound::WavWriter::create("out.wav", spec).unwrap();

for byte in bytes {
for i in 0..8 {
let bit = (byte >> (7 - i)) & 1;
for i in (1..i16::MAX).step_by(32) {
match bit {
0 => {
writer.write_sample(-i).unwrap();
writer.write_sample(i).unwrap();
}
1 => {
writer.write_sample(1).unwrap();
writer.write_sample(1).unwrap();
}
_ => unreachable!(),
}
}
for _ in 0..1 {
match bit {
0 => {
writer.write_sample(i16::MIN).unwrap();
writer.write_sample(i16::MAX).unwrap();
}
1 => {
writer.write_sample(0).unwrap();
writer.write_sample(0).unwrap();
}
_ => unreachable!(),
}
}
for i in (1..i16::MAX).rev().step_by(32) {
match bit {
0 => {
writer.write_sample(-i).unwrap();
writer.write_sample(i).unwrap();
}
1 => {
writer.write_sample(1).unwrap();
writer.write_sample(1).unwrap();
}
_ => unreachable!(),
}
}
}
}
}

pub fn wav_to_bytes(bytes: &[u8]) -> Vec<u8> {
let mut reader = hound::WavReader::new(bytes).unwrap();

let mut bytes = Vec::new();
let mut byte = 0;
let mut bits_read = 0;
// Loop over sets of 4 samples and detect 0 or 1
for chunk in reader
.samples::<i16>()
.collect::<Result<Vec<i16>, _>>()
.unwrap()
.chunks(2)
{
match chunk {
[i16::MIN, i16::MAX] => {
byte <<= 1;
bits_read += 1;
}
[0, 0] => {
byte = (byte << 1) | 1;
bits_read += 1;
}
_ => {}
}

if bits_read == 8 {
bytes.push(byte);
byte = 0;
bits_read = 0;
}
}

bytes
}
13 changes: 7 additions & 6 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
mod cipher;
mod lsb;
mod wav;

use inquire::validator::Validation;
use sneaky::{bytes_to_wav, extract_bit_planes, lsb_decode, lsb_encode, wav_to_bytes};
use std::ops::BitXor;
use std::path::{Path, PathBuf};

Expand Down Expand Up @@ -43,10 +44,10 @@ fn encode() {
let image = image::open(image_path).unwrap().to_rgb8();
let n_bits = get_bits();

lsb_encode(message, image, n_bits, gen_seed());
lsb::encode(message, image, n_bits, gen_seed());
}
"Encode using FSK to a wav" => {
bytes_to_wav(message.as_bytes());
wav::from_bytes(message.as_bytes());
}
"Encode using a ROT cipher" => {
// input any integer
Expand Down Expand Up @@ -89,10 +90,10 @@ fn decode() {
"Decode as an image using LSB" => {
let image = image::load_from_memory(&bytes).unwrap().to_rgb8();
let n_bits = get_bits();
lsb_decode(image, n_bits, gen_seed());
lsb::decode(image, n_bits, gen_seed());
}
"Decode using FSK from a wav" => {
let bytes = wav_to_bytes(&bytes);
let bytes = wav::to_bytes(&bytes);
println!("Message: {}", String::from_utf8_lossy(&bytes));
}
"Decode using a ROT cipher" => {
Expand All @@ -114,7 +115,7 @@ fn analyse() {
"Analyse an image's bit planes" => {
let image_path = get_path();
let image = image::open(image_path).unwrap().to_rgb8();
extract_bit_planes(image);
lsb::extract_bit_planes(image);
}
_ => panic!("Invalid option"),
}
Expand Down
90 changes: 90 additions & 0 deletions src/wav.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
pub(crate) fn from_bytes(bytes: &[u8]) {
let spec = hound::WavSpec {
channels: 1,
sample_rate: 44100,
bits_per_sample: 16,
sample_format: hound::SampleFormat::Int,
};

let mut writer = hound::WavWriter::create("out.wav", spec).unwrap();

for byte in bytes {
for i in 0..8 {
let bit = (byte >> (7 - i)) & 1;
for i in (1..i16::MAX).step_by(64) {
match bit {
0 => {
writer.write_sample(-i).unwrap();
writer.write_sample(i).unwrap();
}
1 => {
writer.write_sample(1).unwrap();
writer.write_sample(1).unwrap();
}
_ => unreachable!(),
}
}
for _ in 0..1 {
match bit {
0 => {
writer.write_sample(i16::MIN).unwrap();
writer.write_sample(i16::MAX).unwrap();
}
1 => {
writer.write_sample(0).unwrap();
writer.write_sample(0).unwrap();
}
_ => unreachable!(),
}
}
for i in (1..i16::MAX).rev().step_by(64) {
match bit {
0 => {
writer.write_sample(-i).unwrap();
writer.write_sample(i).unwrap();
}
1 => {
writer.write_sample(1).unwrap();
writer.write_sample(1).unwrap();
}
_ => unreachable!(),
}
}
}
}
}

pub(crate) fn to_bytes(bytes: &[u8]) -> Vec<u8> {
let mut reader = hound::WavReader::new(bytes).unwrap();

let mut bytes = Vec::new();
let mut byte = 0;
let mut bits_read = 0;
// Loop over sets of 4 samples and detect 0 or 1
for chunk in reader
.samples::<i16>()
.collect::<Result<Vec<i16>, _>>()
.unwrap()
.chunks(2)
{
match chunk {
[i16::MIN, i16::MAX] => {
byte <<= 1;
bits_read += 1;
}
[0, 0] => {
byte = (byte << 1) | 1;
bits_read += 1;
}
_ => {}
}

if bits_read == 8 {
bytes.push(byte);
byte = 0;
bits_read = 0;
}
}

bytes
}

0 comments on commit b51e646

Please sign in to comment.