This commit is contained in:
ef3d0c3e 2024-08-17 13:50:23 +02:00
parent 7e904dc8fc
commit 30d36eb03c
3 changed files with 111 additions and 24 deletions

View file

@ -1,6 +1,11 @@
use std::fmt::Formatter; use std::fmt::Formatter;
use std::str::FromStr; use std::str::FromStr;
use rand::Rng;
use rand::prelude::SliceRandom;
use crate::embed::EmbedAlgorithm;
/// Block mode for embedded data /// Block mode for embedded data
#[derive(Debug)] #[derive(Debug)]
pub struct BlockMode { 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<usize>,
}
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<R>(data: &'a mut [u8], block_size: usize, algorithm: &EmbedAlgorithm, embed_size: usize, rng: &mut R) -> Result<Self, String>
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
})
}
}

View file

@ -12,6 +12,15 @@ pub enum EmbedAlgorithm {
} }
impl 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<dyn ImageInfo>) -> usize { pub fn max_size(&self, blockmode: &BlockMode, info: &Box<dyn ImageInfo>) -> usize {
let blocks_num = info.size() / blockmode.len; 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<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> usize { pub fn next_block(&self, original_data: &mut [u8], mut data_pos: usize, embed_data: &BitVec<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> (usize, usize) {
match self { match self {
EmbedAlgorithm::Lo(bits) => { EmbedAlgorithm::Lo(bits) => {
let mask = (1<<bits) -1; let mask = (1<<bits) -1;
@ -37,23 +46,27 @@ impl EmbedAlgorithm {
result result
} }
for i in 0..(blockmode.len-blockmode.crc_len) let start = embed_offset;
while embed_offset-start < (blockmode.len-blockmode.crc_len)*8
{ {
let hi = std::cmp::min(embed_offset+*bits as usize, embed_data.len())-embed_offset; 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); let embed = bits_to_byte(embed_data.get(embed_offset..embed_offset+hi).unwrap(), hi as u8);
original_data[i] &= !mask; original_data[data_pos] &= !mask;
original_data[i] |= embed; original_data[data_pos] |= embed;
data_pos += 1;
embed_offset += hi; embed_offset += hi;
} }
// TODO: WRITE CRC
} }
} }
embed_offset (data_pos, embed_offset)
} }
pub fn read_block(&self, encoded_data: &[u8], incoming: &mut BitVec<u8>, blockmode: &BlockMode) { pub fn read_block(&self, encoded_data: &[u8], mut data_pos: usize, incoming: &mut BitVec<u8>, blockmode: &BlockMode) -> usize {
match self { match self {
EmbedAlgorithm::Lo(bits) => { EmbedAlgorithm::Lo(bits) => {
fn push(vec: &mut BitVec<u8>, bits: u8, b: u8) fn push(vec: &mut BitVec<u8>, bits: u8, b: u8)
@ -64,17 +77,18 @@ impl EmbedAlgorithm {
} }
} }
let mut i = 0;
let start = incoming.len(); let start = incoming.len();
while incoming.len()-start < (blockmode.len-blockmode.crc_len)*8 while incoming.len()-start < (blockmode.len-blockmode.crc_len)*8
{ {
push(incoming, *bits, encoded_data[i]); push(incoming, *bits, encoded_data[data_pos]);
i += 1; data_pos += 1;
} }
// TODO: Read CRC and verify // TODO: Read CRC and verify
} }
} }
data_pos
} }
} }

View file

@ -13,6 +13,7 @@ use std::str::FromStr;
use bitvec::vec::BitVec; use bitvec::vec::BitVec;
use block::BlockMode; use block::BlockMode;
use block::BlockPlacement;
use embed::EmbedAlgorithm; use embed::EmbedAlgorithm;
use getopts::Matches; use getopts::Matches;
use getopts::Options; use getopts::Options;
@ -121,16 +122,17 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri
// Read header // Read header
let mut read_data = BitVec::<u8>::new(); let mut read_data = BitVec::<u8>::new();
let mut data_pos = 0;
while read_data.len() < 9*8 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()); let (version, blockmode, data_len, comment_len) = Header::from_data(read_data.as_bitslice());
// Read header comment // Read header comment
while read_data.len() < (9+comment_len as usize)*8 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: // Extract comment:
@ -139,9 +141,9 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri
); );
println!("=== HEADER ==="); println!("=== HEADER ===");
println!("Version: {version}"); println!("Version : {version}");
println!("Data Len: {data_len}"); println!("Data Len: {data_len}");
println!("Comment: `{comment}`"); println!("Comment : `{comment}`");
println!("=============="); println!("==============");
fn read_byte(slice: &bitvec::slice::BitSlice<u8>) -> u8 fn read_byte(slice: &bitvec::slice::BitSlice<u8>) -> u8
@ -157,10 +159,10 @@ fn decode(image: String, matches: Matches, header_only: bool) -> Result<(), Stri
let data_start = 9+comment_len as usize; let data_start = 9+comment_len as usize;
while read_data.len() < (data_start + data_len as usize)*8 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]); let b = read_byte(&read_data[(data_start+i)*8..(data_start+1+i)*8]);
println!("{i} : {b:08b} ({})", b as char); println!("{i} : {b:08b} ({})", b as char);
@ -192,11 +194,18 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
.unwrap_or(format!("{}x{}", info.width(), info.height())) .unwrap_or(format!("{}x{}", info.width(), info.height()))
.as_str(), .as_str(),
)?; )?;
let max_size = algorithm.max_size(&blockmode, &info); let max_size = algorithm.max_size(&blockmode, &info);
let embed_data = std::fs::read(&embed_file) let embed_data = std::fs::read(&embed_file)
.map_err(|err| format!("Failed to read embed file `{embed_file}`: {err}"))?; .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 // Get header
let header = Header { let header = Header {
blockmode, blockmode,
@ -252,24 +261,33 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
let mut embed_offset = 0; let mut embed_offset = 0;
let mut data_pos = 0;
for i in 0 .. blocks_num for i in 0 .. blocks_num
{ {
let mut new_offset = algorithm.next_block( println!("block: {i}/{embed_offset}/{data_pos}");
&mut data.as_mut_slice()[i*header.blockmode.len..], (data_pos, embed_offset) = algorithm.next_block(
&mut data.as_mut_slice(),
data_pos,
&bv, &bv,
embed_offset, embed_offset,
&header.blockmode); &header.blockmode);
new_offset += header.blockmode.crc_len*8;
embed_offset = new_offset;
} }
println!("{}", bv.len()); 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); 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>) -> 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 outfile = File::create("out.png").unwrap();
let ref mut w = BufWriter::new(Box::new(outfile) as Box<dyn Write>); let ref mut w = BufWriter::new(Box::new(outfile) as Box<dyn Write>);