From b51e6462d8365d4b729c2fb33cdeabf03fed0fe2 Mon Sep 17 00:00:00 2001 From: Lucas Date: Sun, 5 Nov 2023 00:23:48 +1100 Subject: [PATCH] Refactor lib functions --- src/{lib.rs => lsb.rs} | 109 ++++------------------------------------- src/main.rs | 13 ++--- src/wav.rs | 90 ++++++++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 106 deletions(-) rename src/{lib.rs => lsb.rs} (62%) create mode 100644 src/wav.rs diff --git a/src/lib.rs b/src/lsb.rs similarity index 62% rename from src/lib.rs rename to src/lsb.rs index d59917a..4b6cfd5 100644 --- a/src/lib.rs +++ b/src/lsb.rs @@ -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) { +pub(crate) fn encode(message: String, mut image: image::RgbImage, n_bits: u8, key: Option) { // 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 💀) @@ -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); } } @@ -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 @@ -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) { +pub(crate) fn decode(image: image::RgbImage, n_bits: u8, key: Option) { // Decode the message from the image let message = match key { Some(key) => { @@ -95,9 +95,9 @@ pub fn lsb_decode(image: image::RgbImage, n_bits: u8, key: Option) { 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 @@ -105,7 +105,7 @@ pub fn lsb_decode(image: image::RgbImage, n_bits: u8, key: Option) { println!("Message: {}", message); } -pub fn lsb_raw_decode(carrier: &[u8], n_bits: usize) -> Vec { +fn raw_decode(carrier: &[u8], n_bits: usize) -> Vec { // 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() { @@ -150,7 +150,7 @@ pub fn lsb_raw_decode(carrier: &[u8], n_bits: usize) -> Vec { 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() @@ -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 { - 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::() - .collect::, _>>() - .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 -} diff --git a/src/main.rs b/src/main.rs index 3e9bcc4..f4f5b12 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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}; @@ -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 @@ -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" => { @@ -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"), } diff --git a/src/wav.rs b/src/wav.rs new file mode 100644 index 0000000..06479a8 --- /dev/null +++ b/src/wav.rs @@ -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 { + 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::() + .collect::, _>>() + .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 +}