diff --git a/src/block.rs b/src/block.rs index 2f72dee..9be467d 100644 --- a/src/block.rs +++ b/src/block.rs @@ -1,4 +1,5 @@ -use std::{fmt::Formatter, intrinsics::ctlz, str::FromStr}; +use std::fmt::Formatter; +use std::str::FromStr; /// Block mode for embedded data #[derive(Debug)] @@ -22,18 +23,29 @@ impl BlockMode { Self { len: 1 << best_p, - crc_len: (best_p/4) * crc as usize, + crc_len: (best_p / 4) * crc as usize, } } + pub fn to_data(&self) -> u8 { - ((self.crc_len != 0) as u8) | (ctlz(self.len) << 1) as u8 + ((self.crc_len != 0) as u8) | ((u8::leading_zeros(self.len as u8) + 1) << 1) as u8 + } + + pub fn from_byte(byte: u8) -> BlockMode { + let crc = byte & 0b1; + let len = byte >> 1; + + Self { + len: 1usize << len, + crc_len: (crc * len) as usize + } } } impl core::fmt::Display for BlockMode { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "(len: {}, crc_len: {})", self.len, self.crc_len) - } + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "(len: {}, crc_len: {})", self.len, self.crc_len) + } } impl FromStr for BlockMode { diff --git a/src/embed.rs b/src/embed.rs index 5768e4d..de4679a 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -22,9 +22,7 @@ impl EmbedAlgorithm { } } - pub fn next_block(&self, original_data: &[u8], embed_data: &BitVec, mut embed_offset: usize, blockmode: &BlockMode) -> (Vec, usize) { - let mut result = Vec::::with_capacity(blockmode.len); - + pub fn next_block(&self, original_data: &mut [u8], embed_data: &BitVec, mut embed_offset: usize, blockmode: &BlockMode) -> usize { match self { EmbedAlgorithm::Lo(bits) => { let mask = (1<, blockmode: &BlockMode) { + match self { + EmbedAlgorithm::Lo(bits) => { + fn push(vec: &mut BitVec, bits: u8, b: u8) + { + for i in 0..bits + { + vec.push((b >> i) & 0b1 == 0b1) + } + } + + 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; + } + + // TODO: Read CRC and verify + } + } } } @@ -79,6 +101,7 @@ impl FromStr for EmbedAlgorithm { let value = second.parse::().map_err(|err| { format!("Failed to convert `{second}` to a number of bits: {err}") })?; + // TODO: We can allow more than 8 bits, depending on the image's bit depth if value > 8 || value == 0 { Err(format!( "Cannot specify {value} bits for `lo` method, must be within [1, 8]" diff --git a/src/header.rs b/src/header.rs index 4d0ab60..1189f5e 100644 --- a/src/header.rs +++ b/src/header.rs @@ -1,3 +1,5 @@ +use bitvec::slice::BitSlice; + use crate::block::BlockMode; pub struct Header { @@ -20,7 +22,7 @@ impl Header { header.extend_from_slice(embed_len.to_le_bytes().as_slice()); // Comment len - let comment_len = self.comment.as_ref().map(|c| c.len()).unwrap_or(0); + let comment_len = self.comment.as_ref().map(|c| c.len() as u16).unwrap_or(0 as u16); header.extend_from_slice(comment_len.to_le_bytes().as_slice()); // Comment @@ -30,4 +32,27 @@ impl Header { header } + + pub fn from_data(slice: &BitSlice) -> (u16, BlockMode, u32, u16) { + 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 + } + + let version = ((read_byte(&slice[8..16]) as u16) << 8) | (read_byte(&slice[0..8]) as u16); + let blockmode = BlockMode::from_byte(read_byte(&slice[16..24])); + let len = ((read_byte(&slice[48..56]) as u32) << 24) + | ((read_byte(&slice[40..48]) as u32) << 16) + | ((read_byte(&slice[32..40]) as u32) << 8) + | (read_byte(&slice[24..32]) as u32); + let comment_len = ((read_byte(&slice[64..72]) as u16) << 8) | (read_byte(&slice[56..64]) as u16); + + + (version, blockmode, len, comment_len) + } } diff --git a/src/image.rs b/src/image.rs index e4caaee..5808e09 100644 --- a/src/image.rs +++ b/src/image.rs @@ -1,5 +1,9 @@ +use std::io::BufWriter; +use std::io::Write; + pub trait ImageInfo { fn width(&self) -> u32; fn height(&self) -> u32; fn size(&self) -> usize; + fn encode(&self, w: &mut BufWriter>, data: Vec); } diff --git a/src/main.rs b/src/main.rs index 0d3c200..65076b2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,3 @@ -#![feature(core_intrinsics)] pub mod block; pub mod embed; pub mod header; @@ -6,10 +5,12 @@ pub mod image; use std::env; use std::fs::File; +use std::io::BufWriter; +use std::io::Read; +use std::io::Write; use std::process::ExitCode; use std::str::FromStr; -use bitvec::slice::BitSlice; use bitvec::vec::BitVec; use block::BlockMode; use embed::EmbedAlgorithm; @@ -21,6 +22,7 @@ use rand::SeedableRng; use rand_chacha::ChaCha8Rng; use rand::Rng; use rand::prelude::SliceRandom; +use bitvec::prelude::*; fn print_usage(program: &str, opts: Options) { let brief = format!( @@ -46,6 +48,15 @@ impl ImageInfo for png::OutputInfo { fn height(&self) -> u32 { self.height } fn size(&self) -> usize { self.buffer_size() } + + fn encode(&self, w: &mut BufWriter>, data: Vec) { + let mut encoder = png::Encoder::new(w, self.width(), self.height()); + encoder.set_color(self.color_type); + encoder.set_depth(self.bit_depth); + let mut writer = encoder.write_header().unwrap(); + writer.write_image_data(data.as_slice()).unwrap(); + println!("Ok"); + } } fn get_algorithm(s: Option) -> Result { @@ -93,6 +104,79 @@ fn derive_seed(seed: &str) -> Result<[u8; 32], String> { Ok(result) } +fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), String> { + let algorithm = get_algorithm(matches.opt_str("l"))?; + let crc = false; + + let (data, info) = decode_image(image)?; + let blockmode = BlockMode::from_length(info.size(), crc); + let seed = derive_seed( + matches + .opt_str("s") + .unwrap_or(format!("{}x{}", info.width(), info.height())) + .as_str(), + )?; + + println!("Blockmode: {blockmode}"); + + // Read header + let mut read_data = BitVec::::new(); + while read_data.len() < 9*8 + { + algorithm.read_block(&data[read_data.len()/8..], &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); + } + + // Extract comment: + let comment = String::from_utf8_lossy( + &read_data.as_raw_slice()[9..(9+comment_len as usize)] + ); + + println!("=== HEADER ==="); + println!("Version: {version}"); + println!("Data Len: {data_len}"); + println!("Comment: `{comment}`"); + println!("=============="); + + 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 + } + + 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); + } + + for i in 0..32 + { + let b = read_byte(&read_data[(data_start+i)*8..(data_start+1+i)*8]); + println!("{i} : {b:08b} ({})", b as char); + } + + + + let mut outfile = File::create("decode.png").unwrap(); + outfile.write( + &read_data.as_raw_slice()[data_start..data_start+data_len as usize] + ).unwrap(); + + + Ok(()) +} + fn encode(image: String, matches: Matches) -> Result, String> { let algorithm = get_algorithm(matches.opt_str("l"))?; let crc = false; @@ -100,7 +184,7 @@ fn encode(image: String, matches: Matches) -> Result, String> { .opt_str("i") .ok_or(format!("Embed file is required"))?; - let (data, info) = decode_image(image)?; + let (mut data, info) = decode_image(image)?; let blockmode = BlockMode::from_length(info.size(), crc); let seed = derive_seed( matches @@ -129,24 +213,49 @@ fn encode(image: String, matches: Matches) -> Result, String> { header.blockmode, ))?; } + + // Blocks to write + let blocks_num = ((header_data.len()+embed_data.len()) as f64 / (header.blockmode.len-header.blockmode.crc_len) as f64).ceil() as usize; - // Shuffle the blocks - let mut rand = ChaCha8Rng::from_seed(seed); - let blocks_num = info.size() / (header.blockmode.len-header.blockmode.crc_len); - - let mut blocks_pos = (0..blocks_num).collect::>(); - blocks_pos.shuffle(&mut rand); - - + // Get data let mut bv = BitVec::::from_vec(header_data); bv.extend_from_raw_slice(embed_data.as_slice()); + // zero-padding + while bv.len()/8 < blocks_num*header.blockmode.len + { + for i in 0..8 + { + bv.push(false); + } + } + + // Shuffle the blocks + //let mut rand = ChaCha8Rng::from_seed(seed); + //let mut blocks_pos = (0..blocks_num).collect::>(); + //blocks_pos.shuffle(&mut rand); + + + println!("-------------"); + println!("Writing: {blocks_num}x{} [{}] blocks", header.blockmode.len, header.blockmode.crc_len); + println!("Data length: {} bytes", bv.len()/8); + println!("-------------"); + + //for i in 0..9*4 { + // let b = data[i] & 0b1111; + // println!("{b:b}"); + //} + println!("====="); + + // TODO: make sure the rounding error keep this offset safe + // i.e that two blocks can't overlap + //let coffset = data.len() / (blocks_num+1); + let mut embed_offset = 0; for i in 0 .. blocks_num { - println!("{:#?}", bv.len()); - let (block, mut new_offset) = algorithm.next_block( - &data.as_slice()[i*header.blockmode.len..], + let mut new_offset = algorithm.next_block( + &mut data.as_mut_slice()[i*header.blockmode.len..], &bv, embed_offset, &header.blockmode); @@ -154,10 +263,17 @@ fn encode(image: String, matches: Matches) -> Result, String> { embed_offset = new_offset; } - algorithm.next_block(data.as_slice(), &bv, 48, &header.blockmode); + println!("{}", bv.len()); // TODO: WRITE CRC - println!("Ok"); + + for i in 70..99 { + let b = (data[i*2] & 0b1111) | ((data[i*2+1] & 0b1111) << 4); + println!("{b:08b}, {}", b as char); + } + let outfile = File::create("out.png").unwrap(); + let ref mut w = BufWriter::new(Box::new(outfile) as Box); + info.encode(w, data); Ok(vec![]) } @@ -199,7 +315,15 @@ fn main() -> ExitCode { let input = matches.free[0].clone(); - if matches.opt_present("e") { + if matches.opt_present("z") { + match decode(input, matches, true) { + Ok(_) => todo!(""), + Err(e) => { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } + } else if matches.opt_present("e") { match encode(input, matches) { Ok(_) => todo!(""), Err(e) => {