Working
This commit is contained in:
parent
7e904dc8fc
commit
30d36eb03c
3 changed files with 111 additions and 24 deletions
55
src/block.rs
55
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<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
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
34
src/embed.rs
34
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<dyn ImageInfo>) -> 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<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 {
|
||||
EmbedAlgorithm::Lo(bits) => {
|
||||
let mask = (1<<bits) -1;
|
||||
|
@ -37,23 +46,27 @@ impl EmbedAlgorithm {
|
|||
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);
|
||||
|
||||
original_data[i] &= !mask;
|
||||
original_data[i] |= embed;
|
||||
original_data[data_pos] &= !mask;
|
||||
original_data[data_pos] |= embed;
|
||||
|
||||
data_pos += 1;
|
||||
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 {
|
||||
EmbedAlgorithm::Lo(bits) => {
|
||||
fn push(vec: &mut BitVec<u8>, 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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
46
src/main.rs
46
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::<u8>::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>) -> 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<Vec<u8>, 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<Vec<u8>, 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>) -> 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<dyn Write>);
|
||||
|
|
Loading…
Reference in a new issue