Refactor in progress
This commit is contained in:
parent
30d36eb03c
commit
5535617b74
4 changed files with 315 additions and 49 deletions
228
src/block.rs
228
src/block.rs
|
|
@ -1,8 +1,10 @@
|
|||
use std::fmt::Formatter;
|
||||
use std::str::FromStr;
|
||||
|
||||
use rand::Rng;
|
||||
use bitvec::slice::BitSlice;
|
||||
use bitvec::vec::BitVec;
|
||||
use rand::prelude::SliceRandom;
|
||||
use rand::Rng;
|
||||
|
||||
use crate::embed::EmbedAlgorithm;
|
||||
|
||||
|
|
@ -42,7 +44,7 @@ impl BlockMode {
|
|||
|
||||
Self {
|
||||
len: 1usize << len,
|
||||
crc_len: (crc * len) as usize
|
||||
crc_len: (crc * len) as usize,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -76,50 +78,238 @@ impl FromStr for BlockMode {
|
|||
|
||||
#[derive(Debug)]
|
||||
pub struct BlockPlacement<'a> {
|
||||
algorithm: &'a EmbedAlgorithm,
|
||||
data: &'a mut [u8],
|
||||
block_size: usize,
|
||||
blocks: Vec<usize>,
|
||||
}
|
||||
|
||||
impl<'a> BlockPlacement<'a>
|
||||
{
|
||||
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
|
||||
pub fn new<R>(
|
||||
data: &'a mut [u8],
|
||||
algorithm: &'a EmbedAlgorithm,
|
||||
block_size: usize,
|
||||
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);
|
||||
|
||||
// Maximum number of blocks
|
||||
let max_blocks = data.len() / block_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))
|
||||
if blocks_num > max_blocks {
|
||||
return Err(format!(
|
||||
"Too many blocks required: {blocks_num}, maximum: {max_blocks}"
|
||||
));
|
||||
}
|
||||
|
||||
// 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);
|
||||
}
|
||||
let mut blocks = (0..max_blocks).collect::<Vec<_>>();
|
||||
|
||||
// Shuffle the block order
|
||||
blocks.shuffle(rng);
|
||||
|
||||
// Only keep the first blocks_num blocks
|
||||
blocks.resize(blocks_num, 0);
|
||||
|
||||
Ok(Self {
|
||||
algorithm,
|
||||
data,
|
||||
block_size,
|
||||
blocks
|
||||
blocks,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn write_embed(&mut self)
|
||||
}
|
||||
|
||||
// Iterator over blocks in the resulting image
|
||||
pub struct BlockPlacementIterator<'a> {
|
||||
algorithm: &'a EmbedAlgorithm,
|
||||
data: &'a [u8],
|
||||
block_size: usize,
|
||||
|
||||
// Block index
|
||||
index: usize,
|
||||
// Position of the blocks
|
||||
blocks: Vec<usize>,
|
||||
}
|
||||
|
||||
impl<'a> BlockPlacementIterator<'a> {
|
||||
pub fn new<R: Rng + ?Sized>(
|
||||
algorithm: &'a EmbedAlgorithm,
|
||||
data: &'a [u8],
|
||||
block_size: usize,
|
||||
rng: &mut R,
|
||||
) -> Self {
|
||||
// Maximum number of blocks
|
||||
let max_blocks = data.len() / block_size;
|
||||
|
||||
// Blocks
|
||||
let mut blocks = (0..max_blocks).collect::<Vec<_>>();
|
||||
|
||||
// Shuffle the block order
|
||||
blocks.shuffle(rng);
|
||||
|
||||
Self {
|
||||
algorithm,
|
||||
data,
|
||||
block_size,
|
||||
index: 0,
|
||||
blocks,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Iterator for BlockPlacementIterator<'a> {
|
||||
type Item = Block<'a>;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
if self.index == self.blocks.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
||||
// Block of data in the resulting image
|
||||
#[derive(Debug)]
|
||||
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>,
|
||||
|
||||
// Byte position in [`data`]
|
||||
index: usize,
|
||||
|
||||
// Remainder, i.e. bits that have been read and will be part of the `next` byte
|
||||
remainder: BitVec<u8>,
|
||||
}
|
||||
|
||||
impl<'a> BlockIterator<'a> {
|
||||
pub fn new(block: &'a Block, previous: Option<BlockIterator>) -> Self {
|
||||
if let Some(previous) = previous {
|
||||
Self {
|
||||
block,
|
||||
index: 0,
|
||||
remainder: previous.remainder,
|
||||
}
|
||||
} else {
|
||||
Self {
|
||||
block,
|
||||
index: 0,
|
||||
remainder: BitVec::with_capacity(7),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Iterator for BlockIterator<'a> {
|
||||
type Item = u8;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
let mut byte = 0u8;
|
||||
|
||||
// Read remainder
|
||||
let mut bit_idx = 0;
|
||||
for bit in &self.remainder {
|
||||
byte |= (*bit as u8) << bit_idx;
|
||||
bit_idx += 1;
|
||||
}
|
||||
self.remainder.clear();
|
||||
|
||||
match self.block.0 {
|
||||
EmbedAlgorithm::Lo(bits) => {
|
||||
while bit_idx < 8 {
|
||||
// End of data
|
||||
if self.index == self.block.1.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Read next byte
|
||||
let next = self.block.1[self.index];
|
||||
self.index += 1;
|
||||
|
||||
for i in 0..*bits {
|
||||
if bit_idx < 8 {
|
||||
// Prepend bit to result
|
||||
byte |= ((next >> i) & 0b1) << bit_idx;
|
||||
} else {
|
||||
// Append bit to remainder
|
||||
self.remainder.push((next >> i) & 0b1 == 0b1)
|
||||
}
|
||||
|
||||
bit_idx += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Some(byte)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rand::SeedableRng;
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn block_iterator() {
|
||||
let algorithm = EmbedAlgorithm::Lo(3);
|
||||
let data = vec![
|
||||
0b10111000, 0b11111001, 0b01101010, 0b00111011, 0b11011100, 0b11100110, 0b01100111,
|
||||
0b01100000,
|
||||
];
|
||||
|
||||
let block = Block(&algorithm, &data);
|
||||
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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blockplacement_iterator() {
|
||||
let algorithm = EmbedAlgorithm::Lo(4);
|
||||
let mut 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,
|
||||
];
|
||||
|
||||
let mut rand = ChaCha8Rng::from_seed([1u8; 32]);
|
||||
let mut it = BlockPlacementIterator::new::<_>(&algorithm, &data, 4, &mut rand);
|
||||
|
||||
let mut rand = ChaCha8Rng::from_seed([1u8; 32]);
|
||||
let mut positions = (0..8).collect::<Vec<_>>();
|
||||
positions.shuffle(&mut rand);
|
||||
|
||||
for i in 0..8 {
|
||||
let block = it.next().unwrap();
|
||||
assert_eq!(block.1[0] / 4, positions[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
58
src/embed.rs
58
src/embed.rs
|
|
@ -1,12 +1,13 @@
|
|||
use std::fmt::Formatter;
|
||||
use std::str::FromStr;
|
||||
|
||||
use bitvec::prelude::*;
|
||||
use bitvec::slice::BitSlice;
|
||||
use bitvec::vec::BitVec;
|
||||
|
||||
use crate::block::BlockMode;
|
||||
use crate::image::ImageInfo;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum EmbedAlgorithm {
|
||||
Lo(u8),
|
||||
}
|
||||
|
|
@ -15,9 +16,7 @@ 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
|
||||
}
|
||||
EmbedAlgorithm::Lo(bits) => ((size * 8) as f64 / *bits as f64).ceil() as usize,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -26,35 +25,43 @@ impl EmbedAlgorithm {
|
|||
|
||||
match self {
|
||||
EmbedAlgorithm::Lo(bits) => {
|
||||
(((blockmode.len - blockmode.crc_len)*blocks_num) as f64 * (*bits as f64) / 8f64).floor() as usize
|
||||
(((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<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> (usize, 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;
|
||||
let mask = (1 << bits) - 1;
|
||||
|
||||
fn bits_to_byte(slice: &BitSlice<u8>, bits: u8) -> u8
|
||||
{
|
||||
let mut result : u8 = 0;
|
||||
for i in 0..bits
|
||||
{
|
||||
fn bits_to_byte(slice: &BitSlice<u8>, 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
|
||||
{
|
||||
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);
|
||||
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;
|
||||
}
|
||||
|
|
@ -62,24 +69,27 @@ impl EmbedAlgorithm {
|
|||
// TODO: WRITE CRC
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
(data_pos, embed_offset)
|
||||
}
|
||||
|
||||
pub fn read_block(&self, encoded_data: &[u8], mut data_pos: usize, incoming: &mut BitVec<u8>, blockmode: &BlockMode) -> usize {
|
||||
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)
|
||||
{
|
||||
for i in 0..bits
|
||||
{
|
||||
fn push(vec: &mut BitVec<u8>, 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
|
||||
{
|
||||
while incoming.len() - start < (blockmode.len - blockmode.crc_len) * 8 {
|
||||
push(incoming, *bits, encoded_data[data_pos]);
|
||||
data_pos += 1;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,13 +1,49 @@
|
|||
use bitvec::slice::BitSlice;
|
||||
use bitvec::{slice::BitSlice, vec::BitVec};
|
||||
use crc::Crc;
|
||||
|
||||
use crate::block::BlockMode;
|
||||
|
||||
pub struct Header {
|
||||
pub blockmode: BlockMode,
|
||||
pub comment: Option<String>,
|
||||
#[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<u8>);
|
||||
}
|
||||
|
||||
//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<String>,
|
||||
}
|
||||
|
||||
|
||||
impl Header {
|
||||
pub fn new(version: Version, block_size: usize, data: &[u8], comment: Option<String>) -> 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::<u32>::new(&crc::CRC_32_CKSUM).checksum(data),
|
||||
comment,
|
||||
}
|
||||
}
|
||||
/*
|
||||
pub fn to_data(&self, version: u16, embed_len: u32) -> Vec<u8> {
|
||||
let mut header = vec![];
|
||||
|
||||
|
|
@ -55,4 +91,30 @@ impl Header {
|
|||
|
||||
(version, blockmode, len, comment_len)
|
||||
}
|
||||
*/
|
||||
}
|
||||
|
||||
impl Encode for Header {
|
||||
fn encode(&self, vec: &mut Vec<u8>) {
|
||||
// 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ Public domain\n"
|
|||
);
|
||||
}
|
||||
|
||||
/*
|
||||
impl ImageInfo for png::OutputInfo {
|
||||
fn width(&self) -> u32 { self.width }
|
||||
|
||||
|
|
@ -202,9 +203,9 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
|
|||
.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)?;
|
||||
let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?;
|
||||
|
||||
//return Ok(vec![]);
|
||||
return Ok(vec![]);
|
||||
|
||||
// Get header
|
||||
let header = Header {
|
||||
|
|
@ -295,6 +296,7 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
|
|||
|
||||
Ok(vec![])
|
||||
}
|
||||
*/
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let args: Vec<String> = env::args().collect();
|
||||
|
|
@ -331,6 +333,7 @@ fn main() -> ExitCode {
|
|||
return ExitCode::FAILURE;
|
||||
}
|
||||
|
||||
/*
|
||||
let input = matches.free[0].clone();
|
||||
|
||||
if matches.opt_present("z") {
|
||||
|
|
@ -350,6 +353,7 @@ fn main() -> ExitCode {
|
|||
}
|
||||
}
|
||||
}
|
||||
*/
|
||||
|
||||
ExitCode::SUCCESS
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue