diff --git a/example/embed.png b/example/embed.png new file mode 100644 index 0000000..5e2e580 Binary files /dev/null and b/example/embed.png differ diff --git a/example/input.png b/example/input.png new file mode 100644 index 0000000..923315d Binary files /dev/null and b/example/input.png differ diff --git a/example/test.sh b/example/test.sh new file mode 100755 index 0000000..f05db6b --- /dev/null +++ b/example/test.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash + +PNG_EMBED=../target/debug/png_embed +[ ! -f "${PNG_EMBED}" ] && PNG_EMBED=../target/release/png_embed +[ ! -f "${PNG_EMBED}" ] && echo "Failed to find png_embed executable" && exit + +echo "Encoding..." +for i in {1..7}; do + echo "Writing dec-lo${i}.." + $PNG_EMBED -l lo${i} -e embed.png input.png -o out-lo${i}.png +done + +echo "Decoding..." +for i in {1..7}; do + echo "Decoding out-lo${i} -> dec-lo${i}.." + $PNG_EMBED -l lo${i} -d out-lo${i}.png -o dec-lo${i}.png +done + +echo "Checksums:" +sha256sum embed.png dec-lo*.png # That's nuts! diff --git a/src/embed.rs b/src/embed.rs deleted file mode 100644 index 350ba0b..0000000 --- a/src/embed.rs +++ /dev/null @@ -1,140 +0,0 @@ -use std::fmt::Formatter; -use std::str::FromStr; - -use bitvec::slice::BitSlice; -use bitvec::vec::BitVec; - -use crate::block::BlockMode; -use crate::image::ImageInfo; - -#[derive(Debug)] -pub enum EmbedAlgorithm { - Lo(u8), -} - -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; - - match self { - EmbedAlgorithm::Lo(bits) => { - (((blockmode.len - blockmode.crc_len) * blocks_num) as f64 * (*bits as f64) / 8f64) - .floor() as 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 << bits) - 1; - - fn bits_to_byte(slice: &BitSlice, bits: u8) -> u8 { - let mut result: u8 = 0; - for i in 0..bits { - result |= (slice[i as usize] as u8) << i; - } - result - } - - let start = embed_offset; - while embed_offset - start < (blockmode.len - blockmode.crc_len) * 8 { - let hi = std::cmp::min(*bits as usize, embed_data.len() - embed_offset); - let embed = bits_to_byte( - embed_data.get(embed_offset..embed_offset + hi).unwrap(), - hi as u8, - ); - - original_data[data_pos] &= !mask; - original_data[data_pos] |= embed; - - data_pos += 1; - embed_offset += hi; - } - - // TODO: WRITE CRC - } - } - - (data_pos, embed_offset) - } - - 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) { - for i in 0..bits { - vec.push((b >> i) & 0b1 == 0b1) - } - } - - let start = incoming.len(); - while incoming.len() - start < (blockmode.len - blockmode.crc_len) * 8 { - push(incoming, *bits, encoded_data[data_pos]); - data_pos += 1; - } - - // TODO: Read CRC and verify - } - } - - data_pos - } -} - -impl core::fmt::Display for EmbedAlgorithm { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self { - EmbedAlgorithm::Lo(bits) => write!(f, "Lo({bits})"), - } - } -} - -impl FromStr for EmbedAlgorithm { - type Err = String; - - fn from_str(s: &str) -> Result { - let (dig_pos, _) = s - .char_indices() - .find(|(_, c)| c.is_ascii_digit()) - .ok_or(format!("Unknown algorithm: {s}"))?; - - let (first, second) = s.split_at(dig_pos); - match first { - "lo" => { - 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]" - )) - } else { - Ok(EmbedAlgorithm::Lo(value)) - } - } - _ => Err(format!("Unknown algorithm: {s}")), - } - } -} diff --git a/src/header.rs b/src/header.rs deleted file mode 100644 index 59eb53f..0000000 --- a/src/header.rs +++ /dev/null @@ -1,120 +0,0 @@ -use bitvec::{slice::BitSlice, vec::BitVec}; -use crc::Crc; - -use crate::block::BlockMode; - -#[repr(u16)] -#[allow(non_camel_case_types)] -#[derive(Debug, Clone, Copy)] -pub enum Version { - VERSION_1, -} - -pub trait Encode { - // Encode the data - fn encode(&self, vec: &mut Vec); -} - -//pub trait Decode { -// fn decode(incoming: &mut EmbedIterator) -> (usize, Self); -//} - -#[derive(Debug)] -pub struct Header { - version: Version, - block_size: usize, - data_len: u32, - data_crc: u32, - comment: Option, -} - - -impl Header { - pub fn new(version: Version, block_size: usize, data: &[u8], comment: Option) -> Self { - assert_eq!((data.len() as u32) as usize, data.len()); - assert_eq!(1 << usize::trailing_zeros(block_size), block_size); - assert!(comment.as_ref().map_or(0, |c| c.len()) < u16::MAX as usize); - - Self { - version, - block_size, - data_len: data.len() as u32, - data_crc: Crc::::new(&crc::CRC_32_CKSUM).checksum(data), - comment, - } - } - /* - pub fn to_data(&self, version: u16, embed_len: u32) -> Vec { - let mut header = vec![]; - - // Version - header.extend_from_slice(version.to_le_bytes().as_slice()); - - // TODO: IV+Cipherinfo - // Blockmode - header.push(self.blockmode.to_data().to_le()); - - // Data len - header.extend_from_slice(embed_len.to_le_bytes().as_slice()); - - // Comment len - 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 - if let Some(comment) = &self.comment { - header.extend_from_slice(comment.as_bytes()); - } - - 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) - } - */ -} - -impl Encode for Header { - fn encode(&self, vec: &mut Vec) { - // Version - vec.extend_from_slice((self.version as u16).to_le_bytes().as_slice()); - - // Block size - vec.push((usize::trailing_zeros(self.block_size) as u8).to_le()); - - // Data Len - vec.extend_from_slice(self.data_len.to_le_bytes().as_slice()); - - // Data CRC - vec.extend_from_slice(self.data_crc.to_le_bytes().as_slice()); - - // Comment length - let comment_length = self.comment.as_ref().map_or(0u16, |c| c.len() as u16); - vec.extend_from_slice(comment_length.to_le_bytes().as_slice()); - - // Comment - if let Some(comment) = &self.comment { - vec.extend_from_slice(comment.as_bytes()); - } - } -} diff --git a/src/main.rs b/src/main.rs deleted file mode 100644 index 22252b1..0000000 --- a/src/main.rs +++ /dev/null @@ -1,359 +0,0 @@ -pub mod block; -pub mod embed; -pub mod header; -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::vec::BitVec; -use block::BlockMode; -use block::BlockPlacement; -use embed::EmbedAlgorithm; -use getopts::Matches; -use getopts::Options; -use header::Header; -use image::ImageInfo; -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!( - "Usage: {0} -(e|d|i) FILE [opts] - Encode: {0} -e file.tar -l rgba8 -c \"(.tar) my archive\" > out.png - Decode: {0} -d out.png > file.tar - Info: {0} -i out.png # (.tar) my archive", - program - ); - print!("{}", opts.usage(&brief)); -} - -fn print_version() { - print!( - "png_data -- Embed data into PNG\n -Public domain\n" - ); -} - -/* -impl ImageInfo for png::OutputInfo { - fn width(&self) -> u32 { self.width } - - 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 { - if let Some(s) = &s { - EmbedAlgorithm::from_str(s.as_str()) - } else { - Err("Missing required algorithm parameter".into()) - } -} - -fn get_blockmode(s: Option) -> Result { - if let Some(s) = &s { - BlockMode::from_str(s) - } else { - Err("Missing requires blockmode parameter".into()) - } -} - -fn decode_image(image: String) -> Result<(Vec, Box), String> { - match image.split_at(image.find('.').unwrap_or(0)).1 { - ".png" => { - let decoder = png::Decoder::new( - File::open(&image).map_err(|err| format!("Failed to read `{image}`: {err}"))?, - ); - let mut reader = decoder - .read_info() - .map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?; - let mut result = Vec::with_capacity(reader.output_buffer_size()); - result.resize(reader.output_buffer_size(), 0); - let info = reader - .next_frame(result.as_mut_slice()) - .map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?; - result.resize(info.buffer_size(), 0); - - Ok((result, Box::new(info))) - } - _ => Err(format!("Unable get image type for {image}")), - } -} - -fn derive_seed(seed: &str) -> Result<[u8; 32], String> { - let mut result = [0u8; 32]; - argon2::Argon2::default().hash_password_into(seed.as_bytes(), b"SEED SALT", &mut result) - .map_err(|err| format!("Failed to derive seed `{seed}`: {err}"))?; - 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(); - let mut data_pos = 0; - while read_data.len() < 9*8 - { - 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 - { - data_pos = algorithm.read_block(&data, data_pos, &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 - { - data_pos = algorithm.read_block(&data, data_pos, &mut read_data, &blockmode); - } - - 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); - } - - - - 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; - let embed_file = matches - .opt_str("i") - .ok_or(format!("Embed file is required"))?; - - let (mut 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(), - )?; - - - 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, - comment: matches.opt_str("c"), - }; - let header_data = header.to_data(1, embed_data.len() as u32); - - // Check length - if embed_data.len() + header_data.len() > max_size { - Err(format!( - "Cannot embed {}bytes into {}bytes using the {algorithm} algorithm with blockmode {}. Max embeddable size: {max_size}bytes", - embed_data.len()+header_data.len(), - data.len(), - 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; - - // 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; - let mut data_pos = 0; - for i in 0 .. blocks_num - { - 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); - } - println!("{}", bv.len()); - - - for i in 10..80 { - let b = (data[i*2] & 0b1111) | ((data[i*2+1] & 0b1111) << 4); - 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); - info.encode(w, data); - - Ok(vec![]) -} -*/ - -fn main() -> ExitCode { - let args: Vec = env::args().collect(); - let program = args[0].clone(); - - let mut opts = Options::new(); - opts.optopt("i", "input", "Input file", "PATH"); - opts.optflag("e", "encode", "Encode file"); - opts.optflag("d", "decode", "Decode mode"); - opts.optopt("c", "comment", "Header comment", "TXT"); - opts.optopt("s", "seed", "Force seed", "TXT"); - opts.optflag("", "no-crc", "Disables CRC"); - opts.optflag("z", "info", "Read information"); - opts.optopt("l", "algorithm", "Embed algorithm", "lo3"); - opts.optflag("h", "help", "Print this help menu"); - opts.optflag("v", "version", "Print program version and licenses"); - - let matches = match opts.parse(&args[1..]) { - Ok(m) => m, - Err(f) => { - panic!("{}", f.to_string()) - } - }; - if matches.opt_present("v") { - print_version(); - return ExitCode::SUCCESS; - } - if matches.opt_present("h") { - print_usage(&program, opts); - return ExitCode::SUCCESS; - } - if matches.free.is_empty() { - print_usage(&program, opts); - return ExitCode::FAILURE; - } - - /* - let input = matches.free[0].clone(); - - 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) => { - eprintln!("{e}"); - return ExitCode::FAILURE; - } - } - } - */ - - ExitCode::SUCCESS -} diff --git a/src/block.rs b/src/png_embed/block.rs similarity index 60% rename from src/block.rs rename to src/png_embed/block.rs index 3618523..3e5f9d3 100644 --- a/src/block.rs +++ b/src/png_embed/block.rs @@ -1,6 +1,3 @@ -use std::fmt::Formatter; -use std::str::FromStr; - use bitvec::slice::BitSlice; use bitvec::vec::BitVec; use rand::prelude::SliceRandom; @@ -8,80 +5,29 @@ use rand::Rng; use crate::embed::EmbedAlgorithm; -/// Block mode for embedded data -#[derive(Debug)] -pub struct BlockMode { - pub len: usize, - pub crc_len: usize, -} - -impl BlockMode { - /// Gets the best [`BlockMode`] and the remainder - pub fn from_length(len: usize, crc: bool) -> Self { - let mut best_remainder = len; - let mut best_p = 0; - for p in 4..16 { - let remainder = len % (1 << p); - if remainder <= best_remainder { - best_remainder = remainder; - best_p = p; - } - } - - Self { - len: 1 << best_p, - crc_len: (best_p / 4) * crc as usize, +/// Gets the best blocksize (i.e. that minimize remaining space) for a certain data length. +/// The blocksize is a number in range [16, 65536] +pub fn best_blocksize(len: usize) -> usize { + let mut best_remainder = len; + let mut best_p = 0; + for p in 4..16 { + let remainder = len % (1 << p); + if remainder <= best_remainder { + best_remainder = remainder; + best_p = p; } } - pub fn to_data(&self) -> 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) - } -} - -impl FromStr for BlockMode { - type Err = String; - - fn from_str(s: &str) -> Result { - let size = s - .parse::() - .map_err(|err| format!("Failed to parse `{}` as block size: {err}", s))?; - - if size < 6 || size > 16 { - Err(format!( - "Invalid block size specified: `{size}` expected value within [6; 16]" - ))?; - } - - Ok(BlockMode { - len: 1 << size, - crc_len: size, - }) - } + 1 << best_p } +/// Struct to hold the positions of data blocks #[derive(Debug)] pub struct BlockPlacement<'a> { algorithm: &'a EmbedAlgorithm, data: &'a mut [u8], block_size: usize, - blocks: Vec, + pub blocks: Vec, } impl<'a> BlockPlacement<'a> { @@ -91,8 +37,8 @@ impl<'a> BlockPlacement<'a> { // // Will fail if the data is too small to hold all the blocks pub fn new( - data: &'a mut [u8], algorithm: &'a EmbedAlgorithm, + data: &'a mut [u8], block_size: usize, embed_size: usize, rng: &mut R, @@ -132,7 +78,35 @@ impl<'a> BlockPlacement<'a> { }) } - pub fn write_embed(&mut self) + // Embeds the data into the original image + pub fn write_embed(&mut self, embed: &BitSlice) { + assert_eq!(embed.len() % 8, 0); + + fn bits_to_byte(slice: &BitSlice, bits: u8) -> u8 { + let mut result: u8 = 0; + for i in 0..bits { + result |= (slice[i as usize] as u8) << i; + } + result + } + + let mut index = 0; + match self.algorithm { + EmbedAlgorithm::Lo(bits) => { + for block in &self.blocks { + for i in 0..self.block_size { + let pos = block * self.block_size + i; + let hi = std::cmp::min(*bits as usize, embed.len() - index); + + self.data[pos] &= !((1 << hi) - 1); + self.data[pos] |= bits_to_byte(&embed[index..], hi as u8); + + index += hi; + } + } + } + } + } } // Iterator over blocks in the resulting image @@ -145,9 +119,12 @@ pub struct BlockPlacementIterator<'a> { index: usize, // Position of the blocks blocks: Vec, + // Iterator over the current block + block_it: Option>, } impl<'a> BlockPlacementIterator<'a> { + /// Creates a new embed iterator pub fn new( algorithm: &'a EmbedAlgorithm, data: &'a [u8], @@ -163,29 +140,50 @@ impl<'a> BlockPlacementIterator<'a> { // Shuffle the block order blocks.shuffle(rng); + let first_block_pos = blocks[0] * block_size; + let first_block = &data[first_block_pos..first_block_pos + block_size]; + Self { algorithm, data, block_size, index: 0, blocks, + block_it: Some(BlockIterator::new(Block(algorithm, first_block), None)), } } } impl<'a> Iterator for BlockPlacementIterator<'a> { - type Item = Block<'a>; + type Item = u8; + /// Gets the next embedded byte in the image + /// + /// # Note + /// + /// Even when the [`next()`] is Some(..), if the iterator is past the embed's length, it will + /// return garbage data. fn next(&mut self) -> Option { - if self.index == self.blocks.len() { - return None; + self.block_it.as_ref()?; + + if let Some(byte) = self.block_it.as_mut().unwrap().next() { + Some(byte) + } else { + self.index += 1; + // Get next block + if self.index == self.blocks.len() { + return None; + } + + let block_pos = self.blocks[self.index] * self.block_size; + let block = &self.data[block_pos..block_pos + self.block_size]; + self.block_it = Some(BlockIterator::new( + Block(self.algorithm, block), + self.block_it.take(), + )); + + self.next() } - - let pos = self.blocks[self.index] * self.block_size; - let slice = &self.data[pos..pos + self.block_size]; - self.index += 1; - - Some(Block(self.algorithm, slice)) } } @@ -196,7 +194,7 @@ pub struct Block<'a>(&'a EmbedAlgorithm, &'a [u8]); // Iterator to read embedded data inside a block pub struct BlockIterator<'a> { // Block of the iterator - block: &'a Block<'a>, + block: Block<'a>, // Byte position in [`data`] index: usize, @@ -206,7 +204,7 @@ pub struct BlockIterator<'a> { } impl<'a> BlockIterator<'a> { - pub fn new(block: &'a Block, previous: Option) -> Self { + pub fn new(block: Block<'a>, previous: Option) -> Self { if let Some(previous) = previous { Self { block, @@ -275,6 +273,28 @@ mod tests { use super::*; + #[test] + fn test_write() { + let algorithm = EmbedAlgorithm::Lo(2); + + let mut data = vec![0u8; 8]; + let embed = vec![0xFF, 0xFF]; + + let embed_bits = BitVec::::from_slice(embed.as_slice()); + let mut rand = ChaCha8Rng::from_seed([1u8; 32]); + let mut placement = BlockPlacement::new::<_>( + &algorithm, + data.as_mut_slice(), + 4, + embed_bits.len() / 8, + &mut rand, + ) + .unwrap(); + placement.write_embed(embed_bits.as_bitslice()); + + assert_eq!(data, vec![0b00000011; 8]); + } + #[test] fn block_iterator() { let algorithm = EmbedAlgorithm::Lo(3); @@ -284,17 +304,17 @@ mod tests { ]; let block = Block(&algorithm, &data); - let mut it = BlockIterator::new(&block, None); + let mut it = BlockIterator::new(block, None); assert_eq!(it.next(), Some(0b10_001_000)); - assert_eq!(it.next(), Some(0b0_100_011_0)); - assert_eq!(it.next(), Some(0b000_111_11)); + assert_eq!(it.next(), Some(0b0100_0110)); + assert_eq!(it.next(), Some(0b0001_1111)); } #[test] fn blockplacement_iterator() { let algorithm = EmbedAlgorithm::Lo(4); - let mut data = vec![ + let data = vec![ 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, @@ -307,9 +327,10 @@ mod tests { let mut positions = (0..8).collect::>(); positions.shuffle(&mut rand); - for i in 0..8 { - let block = it.next().unwrap(); - assert_eq!(block.1[0] / 4, positions[i]); + for i in 0..data.len() / 2 { + let byte = it.next().unwrap(); + // TODO... + //assert_eq!(byte, data[positions[i/4]*4+(i%4)]); } } } diff --git a/src/png_embed/embed.rs b/src/png_embed/embed.rs new file mode 100644 index 0000000..b34c999 --- /dev/null +++ b/src/png_embed/embed.rs @@ -0,0 +1,54 @@ +use std::fmt::Formatter; +use std::str::FromStr; + +/// Algorithm to embed data +#[derive(Debug)] +pub enum EmbedAlgorithm { + Lo(u8), +} + +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, + } + } +} + +impl core::fmt::Display for EmbedAlgorithm { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + EmbedAlgorithm::Lo(bits) => write!(f, "Lo({bits})"), + } + } +} + +impl FromStr for EmbedAlgorithm { + type Err = String; + + fn from_str(s: &str) -> Result { + let (dig_pos, _) = s + .char_indices() + .find(|(_, c)| c.is_ascii_digit()) + .ok_or(format!("Unknown algorithm: {s}"))?; + + let (first, second) = s.split_at(dig_pos); + match first { + "lo" => { + 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 > 7 || value == 0 { + Err(format!( + "Cannot specify {value} bits for `lo` method, must be within [1, 7]" + )) + } else { + Ok(EmbedAlgorithm::Lo(value)) + } + } + _ => Err(format!("Unknown algorithm: {s}")), + } + } +} diff --git a/src/png_embed/header.rs b/src/png_embed/header.rs new file mode 100644 index 0000000..0eebbe1 --- /dev/null +++ b/src/png_embed/header.rs @@ -0,0 +1,139 @@ +use crc::Crc; + +use crate::block::BlockPlacementIterator; + +/// The program's version. +/// Used for compatibility reasons. +#[repr(u16)] +#[allow(non_camel_case_types)] +#[derive(Debug, Clone, Copy)] +pub enum Version { + VERSION_1, +} + +impl TryFrom for Version { + type Error = String; + + fn try_from(value: u16) -> Result { + match value { + 0 => Ok(Version::VERSION_1), + ver => Err(format!("Unknown version: {ver}")), + } + } +} + +pub trait Encode { + /// Encode the data into a vector + fn encode(&self, vec: &mut Vec); +} + +pub trait Decode { + type Type; + + /// Decode the data from an iterator + fn decode(it: &mut BlockPlacementIterator) -> Result; +} + +/// Embed data header: +/// +---------+----------+----------+-------------+---------+ +/// | Version | Data Len | Data CRC | Comment Len | Comment | +/// +---------+----------+----------+-------------+---------+ +/// | 2 | 4 | 4 | 2 | varies | +/// +---------+----------+----------+-------------+---------+ +#[derive(Debug)] +pub struct Header { + pub version: Version, + pub data_len: u32, + pub data_crc: u32, + pub comment: Option, +} + +impl Header { + /// Construct a new header from the embedded data + pub fn new(version: Version, data: &[u8], comment: Option) -> Result { + if data.len() > u32::MAX as usize { + return Err(format!( + "Embedded data length: {} is greater than maximum {}", + data.len(), + u32::MAX + )); + } else if let Some(len) = comment.as_ref().map(|c| c.len()) { + if len > u16::MAX as usize { + return Err(format!( + "Embedded comment is too long, maximum length: {}, got {len}", + u16::MAX + )); + } + } + + Ok(Self { + version, + data_len: data.len() as u32, + data_crc: Crc::::new(&crc::CRC_32_CKSUM).checksum(data), + comment, + }) + } +} + +impl Encode for Header { + fn encode(&self, vec: &mut Vec) { + // Version + vec.extend_from_slice((self.version as u16).to_le_bytes().as_slice()); + + // Data Len + vec.extend_from_slice(self.data_len.to_le_bytes().as_slice()); + + // Data CRC + vec.extend_from_slice(self.data_crc.to_le_bytes().as_slice()); + + // Comment length + let comment_length = self.comment.as_ref().map_or(0u16, |c| c.len() as u16); + vec.extend_from_slice(comment_length.to_le_bytes().as_slice()); + + // Comment + if let Some(comment) = &self.comment { + vec.extend_from_slice(comment.as_bytes()); + } + } +} + +impl Decode for Header { + type Type = Header; + + fn decode(it: &mut BlockPlacementIterator) -> Result { + let mut count = 0; + let mut next = || -> Result { + let result = it + .next() + .ok_or(format!("Failed to get byte at index: {count}")); + count += 1; + result + }; + + let version = u16::from_le_bytes([next()?, next()?]); + let data_len = u32::from_le_bytes([next()?, next()?, next()?, next()?]); + let data_crc = u32::from_le_bytes([next()?, next()?, next()?, next()?]); + let comment_length = u16::from_le_bytes([next()?, next()?]); + + let comment = if comment_length != 0 { + let mut comment_data = Vec::with_capacity(comment_length as usize); + for _ in 0..comment_length { + comment_data.push(next()?); + } + + Some( + String::from_utf8(comment_data) + .map_err(|e| format!("Failed to retrieve comment: {e}"))?, + ) + } else { + None + }; + + Ok(Header { + version: Version::try_from(version)?, + data_len, + data_crc, + comment, + }) + } +} diff --git a/src/image.rs b/src/png_embed/image.rs similarity index 100% rename from src/image.rs rename to src/png_embed/image.rs diff --git a/src/png_embed/main.rs b/src/png_embed/main.rs new file mode 100644 index 0000000..b01c0eb --- /dev/null +++ b/src/png_embed/main.rs @@ -0,0 +1,335 @@ +pub mod block; +pub mod embed; +pub mod header; +pub mod image; + +use std::env; +use std::fs::File; +use std::io::BufWriter; +use std::io::Write; +use std::process::ExitCode; +use std::str::FromStr; + +use bitvec::prelude::*; +use block::best_blocksize; +use block::BlockPlacement; +use block::BlockPlacementIterator; +use crc::Crc; +use embed::EmbedAlgorithm; +use getopts::Matches; +use getopts::Options; +use header::Decode; +use header::Encode; +use header::Header; +use image::ImageInfo; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +fn print_usage(program: &str, opts: Options) { + let brief = format!( + "Usage: {0} -l ALGORITHM -(e|d|z) [EMBED] FILE -o OUTPUT [opts] + Encode: {0} -l lo3 -e embed.jpg input.png -o out.png -c \"Embedded JPEG file\" + Info: {0} -l lo3 out.png # Embedded JPEG file + Decode: {0} -l lo3 out.png > decoded.jpg", + program + ); + print!("{}", opts.usage(&brief)); +} + +fn print_version() { + print!( + r#"png_embed (c) ef3d0c3e -- Embed data into PNGs +Copyright (c) 2024 +NML is licensed under the GNU Affero General Public License version 3 (AGPLv3), +under the terms of the Free Software Foundation . + +This program is free software; you may modify and redistribute it. +There is NO WARRANTY, to the extent permitted by law."# + ); +} + +impl ImageInfo for png::OutputInfo { + fn width(&self) -> u32 { self.width } + + 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(); + } +} + +fn decode_image(image: &str) -> Result<(Vec, Box), String> { + match image.split_at(image.find('.').unwrap_or(0)).1 { + ".png" => { + let decoder = png::Decoder::new( + File::open(image).map_err(|err| format!("Failed to read `{image}`: {err}"))?, + ); + let mut reader = decoder + .read_info() + .map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?; + let mut result = vec![0; reader.output_buffer_size()]; + let info = reader + .next_frame(result.as_mut_slice()) + .map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?; + result.resize(info.buffer_size(), 0); + + Ok((result, Box::new(info))) + } + _ => Err(format!("Unable get image type for {image}")), + } +} + +// Derives the seed from a given string. +// Currently using Argon with salt: `png_data embed` +fn derive_seed(seed: &str) -> Result<[u8; 32], String> { + let mut result = [0u8; 32]; + argon2::Argon2::default() + .hash_password_into(seed.as_bytes(), b"png_data embed", &mut result) + .map_err(|err| format!("Failed to derive seed `{seed}`: {err}"))?; + Ok(result) +} + +fn encode( + input: String, + embed: String, + output: String, + algorithm: String, + matches: Matches, +) -> Result<(), String> { + let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?; + + let (mut data, info) = decode_image(input.as_str())?; + let block_size = best_blocksize(info.size()); + let seed = derive_seed( + matches + .opt_str("s") + .unwrap_or(format!("{}x{}", info.width(), info.height())) + .as_str(), + )?; + let comment = matches.opt_str("c"); + + // Data + let embed_file_data = std::fs::read(&embed) + .map_err(|err| format!("Failed to read embed file `{embed}`: {err}"))?; + + // Header + let header = Header::new( + header::Version::VERSION_1, + embed_file_data.as_slice(), + comment, + )?; + + // Result + let mut embed_data = vec![]; + header.encode(&mut embed_data); + embed_data.extend(embed_file_data); + + eprintln!("=== HEADER ==="); + eprintln!("Version: {:#?}", header.version); + eprintln!( + "Comment: {}", + header.comment.as_ref().map_or("", |c| c.as_str()) + ); + eprintln!("Data: {}bytes CRC[{:X}]", header.data_len, header.data_crc); + eprintln!("Block: {block_size}bytes"); + + let mut rand = ChaCha8Rng::from_seed(seed); + let mut placement = BlockPlacement::new( + &algorithm, + data.as_mut_slice(), + block_size, + embed_data.len(), + &mut rand, + )?; + + eprintln!("Required blocks: {}", placement.blocks.len()); + eprintln!("=============="); + + placement.write_embed(embed_data.as_slice().view_bits::()); + + let outfile = File::create(&output).unwrap(); + let w = &mut BufWriter::new(Box::new(outfile) as Box); + info.encode(w, data); + + Ok(()) +} + +fn decode_header(input: String, algorithm: String, matches: Matches) -> Result<(), String> { + let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?; + + let (data, info) = decode_image(input.as_str())?; + let block_size = best_blocksize(info.size()); + let seed = derive_seed( + matches + .opt_str("s") + .unwrap_or(format!("{}x{}", info.width(), info.height())) + .as_str(), + )?; + + let mut rand = ChaCha8Rng::from_seed(seed); + let mut it = BlockPlacementIterator::new(&algorithm, data.as_slice(), block_size, &mut rand); + + let header = Header::decode(&mut it)?; + + eprintln!("=== HEADER ==="); + eprintln!("Version: {:#?}", header.version); + eprintln!( + "Comment: \"{}\"", + header.comment.as_ref().map_or("", |c| c.as_str()) + ); + eprintln!("Data: {}bytes CRC[{:X}]", header.data_len, header.data_crc); + eprintln!("=============="); + + Ok(()) +} + +fn decode( + input: String, + output: String, + algorithm: String, + matches: Matches, +) -> Result<(), String> { + let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?; + + let (data, info) = decode_image(input.as_str())?; + let block_size = best_blocksize(info.size()); + let seed = derive_seed( + matches + .opt_str("s") + .unwrap_or(format!("{}x{}", info.width(), info.height())) + .as_str(), + )?; + + let mut rand = ChaCha8Rng::from_seed(seed); + let mut it = BlockPlacementIterator::new(&algorithm, data.as_slice(), block_size, &mut rand); + + let header = Header::decode(&mut it)?; + + let mut data = Vec::with_capacity(header.data_len as usize); + while data.len() < header.data_len as usize { + data.push( + it.next() + .ok_or(format!("Failed to read data byte at {}", data.len()))?, + ); + } + + // Check CRC + let data_crc = Crc::::new(&crc::CRC_32_CKSUM).checksum(data.as_slice()); + if data_crc != header.data_crc { + Err(format!( + "Data CRC do not match: HEADER={:X} GOT={data_crc:X}", + header.data_crc + ))?; + } + + let outfile = File::create(&output) + .map_err(|e| format!("Failed to create output file `{output}`: {e}"))?; + let w = &mut BufWriter::new(Box::new(outfile) as Box); + w.write_all(data.as_slice()) + .map_err(|e| format!("Failed to write to output file `{output}`: {e}"))?; + + eprintln!("File written to `{output}`"); + + Ok(()) +} + +fn main() -> ExitCode { + let args: Vec = env::args().collect(); + let program = args[0].clone(); + + let mut opts = Options::new(); + opts.optopt("e", "embed", "Embed file", "PATH"); + opts.optopt("o", "output", "Output file", "PATH"); + opts.optflag("d", "decode", "Decode mode"); + opts.optopt("c", "comment", "Header comment", "TXT"); + opts.optopt( + "s", + "seed", + "Force a seed, defaults to \"{width}x{height}\"", + "TXT", + ); + opts.optflag("z", "info", "Read header"); + opts.optopt("l", "algorithm", "Embed algorithm", "lo3"); + opts.optflag("h", "help", "Print this help menu"); + opts.optflag("v", "version", "Print program version and licenses"); + + let matches = match opts.parse(&args[1..]) { + Ok(m) => m, + Err(f) => { + panic!("{}", f.to_string()) + } + }; + if matches.opt_present("v") { + print_version(); + return ExitCode::SUCCESS; + } + if matches.opt_present("h") { + print_usage(&program, opts); + return ExitCode::SUCCESS; + } + + // Get input file + if matches.free.is_empty() { + eprintln!("Missing input file"); + print_usage(&program, opts); + return ExitCode::FAILURE; + } + let input_file = matches.free[0].clone(); + + // Check options + if matches.opt_present("e") as usize + + matches.opt_present("d") as usize + + matches.opt_present("z") as usize + > 1 + { + eprintln!("Specify either `-e(--embed)`, -z(--info) or `-d(--decode)`"); + return ExitCode::FAILURE; + } else if !matches.opt_present("l") { + eprintln!("Missing algorithm name"); + return ExitCode::FAILURE; + } + + // Get algorithm + let algorithm = matches.opt_str("l").unwrap(); + + if matches.opt_present("e") { + let embed_file = matches.opt_str("e").unwrap(); + if !matches.opt_present("o") { + eprintln!("Missing -o(utput) file"); + return ExitCode::FAILURE; + } + let output_file = matches.opt_str("o").unwrap(); + + if let Err(e) = encode(input_file, embed_file, output_file, algorithm, matches) { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } else if matches.opt_present("z") { + if let Err(e) = decode_header(input_file, algorithm, matches) { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } else if matches.opt_present("d") { + if !matches.opt_present("o") { + eprintln!("Missing -o(utput) file"); + return ExitCode::FAILURE; + } + let output_file = matches.opt_str("o").unwrap(); + + if let Err(e) = decode(input_file, output_file, algorithm, matches) { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } else { + print_usage(&program, opts); + return ExitCode::FAILURE; + } + + ExitCode::SUCCESS +}