Progress

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,11 +23,22 @@ 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
+		}
 	}
 }
 

src/embed.rs

@@ -22,9 +22,7 @@ impl EmbedAlgorithm {
 		}
 	}
 
-	pub fn next_block(&self, original_data: &[u8], embed_data: &BitVec<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> (Vec<u8>, usize) {
-		let mut result = Vec::<u8>::with_capacity(blockmode.len);
-
+	pub fn next_block(&self, original_data: &mut [u8], embed_data: &BitVec<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> usize {
 		match self {
 			EmbedAlgorithm::Lo(bits) => {
 				let mask = (1<<bits) -1;
@@ -41,18 +39,42 @@ impl EmbedAlgorithm {
 
 				for i in 0..(blockmode.len-blockmode.crc_len)
 				{
-					let embed = bits_to_byte(embed_data.get(embed_offset..embed_offset+*bits as usize).unwrap(), *bits);
-					let b = original_data[i];
+					let hi = std::cmp::min(embed_offset+*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);
 
-					result.push((b & !mask) | embed);
+					original_data[i] &= !mask;
+					original_data[i] |= embed;
 					
-					embed_offset += *bits as usize;
+					embed_offset += hi;
 				}
 			}
 		}
 		
+		embed_offset
+	}
 
-		(result, embed_offset)
+	pub fn read_block(&self, encoded_data: &[u8], incoming: &mut BitVec<u8>, blockmode: &BlockMode) {
+		match self {
+			EmbedAlgorithm::Lo(bits) => {
+				fn push(vec: &mut BitVec<u8>, 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::<u8>().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]"

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<u8>) -> (u16, BlockMode, u32, u16) {
+		fn read_byte(slice: &bitvec::slice::BitSlice<u8>) -> 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)
+	}
 }

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<Box<dyn Write>>, data: Vec<u8>);
 }

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<Box<dyn Write>>, data: Vec<u8>) {
+		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<String>) -> Result<EmbedAlgorithm, String> {
@@ -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::<u8>::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>) -> 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<Vec<u8>, String> {
 	let algorithm = get_algorithm(matches.opt_str("l"))?;
 	let crc = false;
@@ -100,7 +184,7 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, 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
@@ -130,23 +214,48 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
 		))?;
 	}
 	
-	// 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::<Vec<_>>();
-	blocks_pos.shuffle(&mut rand);
-
+	// 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::<u8>::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::<Vec<_>>();
+	//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<Vec<u8>, 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<dyn Write>);
+	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) => {