ef3d0c3e/png_data
3
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Progress

Browse source
This commit is contained in:
ef3d0c3e 2024-08-16 22:51:49 +02:00
parent f9728d490d
commit 7e904dc8fc
5 changed files with 220 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
src/block.rs
View file

@ -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,18 +23,29 @@ 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
}
}
}
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)
}
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "(len: {}, crc_len: {})", self.len, self.crc_len)
}
}
impl FromStr for BlockMode {

39
src/embed.rs
View file

@ -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]"

27
src/header.rs
View file

@ -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)
}
}

4
src/image.rs
View file

@ -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>);
}

158
src/main.rs
View file

@ -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
@ -129,24 +213,49 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
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;
// 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);
// 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) => {