diff --git a/src/block.rs b/src/block.rs index 9be467d..977bb44 100644 --- a/src/block.rs +++ b/src/block.rs @@ -1,6 +1,11 @@ use std::fmt::Formatter; use std::str::FromStr; +use rand::Rng; +use rand::prelude::SliceRandom; + +use crate::embed::EmbedAlgorithm; + /// Block mode for embedded data #[derive(Debug)] pub struct BlockMode { @@ -68,3 +73,53 @@ impl FromStr for BlockMode { }) } } + +#[derive(Debug)] +pub struct BlockPlacement<'a> { + data: &'a mut [u8], + block_size: usize, + blocks: Vec, +} + +impl<'a> BlockPlacement<'a> +{ + // Attempts to create a new block placement + // + // # Errors + // + // Will fail if the data is too small to hold all the blocks + pub fn new(data: &'a mut [u8], block_size: usize, algorithm: &EmbedAlgorithm, embed_size: usize, rng: &mut R) -> Result + where R: Rng + ?Sized + { + // Total size of the embed (crc included) + let embedded_size = algorithm.embedded_size(embed_size); + + // Number of blocks + let blocks_num = (embedded_size as f64 / block_size as f64).ceil() as usize; + + + // Safe distance for spacing the blocks equally + let safe_spacing = data.len() / blocks_num; + if safe_spacing*blocks_num < embedded_size { + return Err(format!("Failed to determine a safe spacing size: {safe_spacing} < {}", embedded_size as f64 / blocks_num as f64)) + } + + // Blocks in the resulting data + let mut blocks = Vec::with_capacity(blocks_num); + for i in 0 .. blocks_num { + // Choose a random position within [0, safe_spacing[ for the block + let pos = rng.gen_range(i*safe_spacing..(i+1)*safe_spacing); + + blocks.push(pos); + } + + // Shuffle the block order + blocks.shuffle(rng); + + Ok(Self { + data, + block_size, + blocks + }) + } +} diff --git a/src/embed.rs b/src/embed.rs index de4679a..fa3afd9 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -12,6 +12,15 @@ pub enum EmbedAlgorithm { } impl EmbedAlgorithm { + /// Get the size of the data (in bytes) once embedded by the algorithm + pub fn embedded_size(&self, size: usize) -> usize { + match self { + EmbedAlgorithm::Lo(bits) => { + ((size * 8) as f64 / *bits as f64).ceil() as usize + } + } + } + pub fn max_size(&self, blockmode: &BlockMode, info: &Box) -> usize { let blocks_num = info.size() / blockmode.len; @@ -22,7 +31,7 @@ impl EmbedAlgorithm { } } - pub fn next_block(&self, original_data: &mut [u8], embed_data: &BitVec, mut embed_offset: usize, blockmode: &BlockMode) -> usize { + pub fn next_block(&self, original_data: &mut [u8], mut data_pos: usize, embed_data: &BitVec, mut embed_offset: usize, blockmode: &BlockMode) -> (usize, usize) { match self { EmbedAlgorithm::Lo(bits) => { let mask = (1<, blockmode: &BlockMode) { + pub fn read_block(&self, encoded_data: &[u8], mut data_pos: usize, incoming: &mut BitVec, blockmode: &BlockMode) -> usize { match self { EmbedAlgorithm::Lo(bits) => { fn push(vec: &mut BitVec, bits: u8, b: u8) @@ -64,17 +77,18 @@ impl EmbedAlgorithm { } } - let mut i = 0; let start = incoming.len(); while incoming.len()-start < (blockmode.len-blockmode.crc_len)*8 { - push(incoming, *bits, encoded_data[i]); - i += 1; + push(incoming, *bits, encoded_data[data_pos]); + data_pos += 1; } // TODO: Read CRC and verify } } + + data_pos } } diff --git a/src/main.rs b/src/main.rs index 65076b2..56666ac 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,6 +13,7 @@ use std::str::FromStr; use bitvec::vec::BitVec; use block::BlockMode; +use block::BlockPlacement; use embed::EmbedAlgorithm; use getopts::Matches; use getopts::Options; @@ -121,16 +122,17 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri // Read header let mut read_data = BitVec::::new(); + let mut data_pos = 0; while read_data.len() < 9*8 { - algorithm.read_block(&data[read_data.len()/8..], &mut read_data, &blockmode); + data_pos = algorithm.read_block(&data, data_pos, &mut read_data, &blockmode); } let (version, blockmode, data_len, comment_len) = Header::from_data(read_data.as_bitslice()); // Read header comment while read_data.len() < (9+comment_len as usize)*8 { - algorithm.read_block(&data[read_data.len()/8..], &mut read_data, &blockmode); + data_pos = algorithm.read_block(&data, data_pos, &mut read_data, &blockmode); } // Extract comment: @@ -139,9 +141,9 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri ); println!("=== HEADER ==="); - println!("Version: {version}"); + println!("Version : {version}"); println!("Data Len: {data_len}"); - println!("Comment: `{comment}`"); + println!("Comment : `{comment}`"); println!("=============="); fn read_byte(slice: &bitvec::slice::BitSlice) -> u8 @@ -157,10 +159,10 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri let data_start = 9+comment_len as usize; while read_data.len() < (data_start + data_len as usize)*8 { - algorithm.read_block(&data[read_data.len()/8..], &mut read_data, &blockmode); + data_pos = algorithm.read_block(&data, data_pos, &mut read_data, &blockmode); } - for i in 0..32 + for i in 60..80 { let b = read_byte(&read_data[(data_start+i)*8..(data_start+1+i)*8]); println!("{i} : {b:08b} ({})", b as char); @@ -192,11 +194,18 @@ fn encode(image: String, matches: Matches) -> Result, String> { .unwrap_or(format!("{}x{}", info.width(), info.height())) .as_str(), )?; + + let max_size = algorithm.max_size(&blockmode, &info); let embed_data = std::fs::read(&embed_file) .map_err(|err| format!("Failed to read embed file `{embed_file}`: {err}"))?; + let mut rand = ChaCha8Rng::from_seed(seed); + //let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?; + + //return Ok(vec![]); + // Get header let header = Header { blockmode, @@ -252,24 +261,33 @@ fn encode(image: String, matches: Matches) -> Result, String> { let mut embed_offset = 0; + let mut data_pos = 0; for i in 0 .. blocks_num { - let mut new_offset = algorithm.next_block( - &mut data.as_mut_slice()[i*header.blockmode.len..], + println!("block: {i}/{embed_offset}/{data_pos}"); + (data_pos, embed_offset) = algorithm.next_block( + &mut data.as_mut_slice(), + data_pos, &bv, embed_offset, &header.blockmode); - new_offset += header.blockmode.crc_len*8; - - embed_offset = new_offset; } println!("{}", bv.len()); - // TODO: WRITE CRC - for i in 70..99 { + for i in 10..80 { let b = (data[i*2] & 0b1111) | ((data[i*2+1] & 0b1111) << 4); - println!("{b:08b}, {}", b as char); + println!("{i}: {b:08b}, {}", b as char); + fn read_byte(slice: &bitvec::slice::BitSlice) -> u8 + { + let mut result = 0; + for i in 0..8 + { + result |= (slice[i as usize] as u8) << i; + } + result + } + println!("{i}+ {b:08b}, {}", read_byte(&bv[i*8..(i+1)*8]) as char); } let outfile = File::create("out.png").unwrap(); let ref mut w = BufWriter::new(Box::new(outfile) as Box);