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

Refactored

Browse source
This commit is contained in:
ef3d0c3e 2024-08-18 16:02:16 +02:00
parent 5535617b74
commit ba9e61ace1
11 changed files with 656 additions and 706 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
example/embed.png Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 65 KiB

BIN
example/input.png Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 64 KiB

20
example/test.sh Executable file
View file

@ -0,0 +1,20 @@
#!/usr/bin/env bash
PNG_EMBED=../target/debug/png_embed
[ ! -f "${PNG_EMBED}" ] && PNG_EMBED=../target/release/png_embed
[ ! -f "${PNG_EMBED}" ] && echo "Failed to find png_embed executable" && exit
echo "Encoding..."
for i in {1..7}; do
echo "Writing dec-lo${i}.."
$PNG_EMBED -l lo${i} -e embed.png input.png -o out-lo${i}.png
done
echo "Decoding..."
for i in {1..7}; do
echo "Decoding out-lo${i} -> dec-lo${i}.."
$PNG_EMBED -l lo${i} -d out-lo${i}.png -o dec-lo${i}.png
done
echo "Checksums:"
sha256sum embed.png dec-lo*.png # That's nuts!

140
src/embed.rs
View file

@ -1,140 +0,0 @@
use std::fmt::Formatter;
use std::str::FromStr;
use bitvec::slice::BitSlice;
use bitvec::vec::BitVec;
use crate::block::BlockMode;
use crate::image::ImageInfo;
#[derive(Debug)]
pub enum EmbedAlgorithm {
Lo(u8),
}
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;
match self {
EmbedAlgorithm::Lo(bits) => {
(((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) {
match self {
EmbedAlgorithm::Lo(bits) => {
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 {
result |= (slice[i as usize] as u8) << i;
}
result
}
let start = embed_offset;
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,
);
original_data[data_pos] &= !mask;
original_data[data_pos] |= embed;
data_pos += 1;
embed_offset += hi;
}
// 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 {
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 start = incoming.len();
while incoming.len() - start < (blockmode.len - blockmode.crc_len) * 8 {
push(incoming, *bits, encoded_data[data_pos]);
data_pos += 1;
}
// TODO: Read CRC and verify
}
}
data_pos
}
}
impl core::fmt::Display for EmbedAlgorithm {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
EmbedAlgorithm::Lo(bits) => write!(f, "Lo({bits})"),
}
}
}
impl FromStr for EmbedAlgorithm {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let (dig_pos, _) = s
.char_indices()
.find(|(_, c)| c.is_ascii_digit())
.ok_or(format!("Unknown algorithm: {s}"))?;
let (first, second) = s.split_at(dig_pos);
match first {
"lo" => {
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]"
))
} else {
Ok(EmbedAlgorithm::Lo(value))
}
}
_ => Err(format!("Unknown algorithm: {s}")),
}
}
}

120
src/header.rs
View file

@ -1,120 +0,0 @@
use bitvec::{slice::BitSlice, vec::BitVec};
use crc::Crc;
use crate::block::BlockMode;
#[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![];
// Version
header.extend_from_slice(version.to_le_bytes().as_slice());
// TODO: IV+Cipherinfo
// Blockmode
header.push(self.blockmode.to_data().to_le());
// Data len
header.extend_from_slice(embed_len.to_le_bytes().as_slice());
// Comment len
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
if let Some(comment) = &self.comment {
header.extend_from_slice(comment.as_bytes());
}
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)
}
*/
}
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());
}
}
}

359
src/main.rs
View file

@ -1,359 +0,0 @@
pub mod block;
pub mod embed;
pub mod header;
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::vec::BitVec;
use block::BlockMode;
use block::BlockPlacement;
use embed::EmbedAlgorithm;
use getopts::Matches;
use getopts::Options;
use header::Header;
use image::ImageInfo;
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!(
"Usage: {0} -(e|d|i) FILE [opts]
Encode: {0} -e file.tar -l rgba8 -c \"(.tar) my archive\" > out.png
Decode: {0} -d out.png > file.tar
Info: {0} -i out.png # (.tar) my archive",
program
);
print!("{}", opts.usage(&brief));
}
fn print_version() {
print!(
"png_data -- Embed data into PNG\n
Public domain\n"
);
}
/*
impl ImageInfo for png::OutputInfo {
fn width(&self) -> u32 { self.width }
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> {
if let Some(s) = &s {
EmbedAlgorithm::from_str(s.as_str())
} else {
Err("Missing required algorithm parameter".into())
}
}
fn get_blockmode(s: Option<String>) -> Result<BlockMode, String> {
if let Some(s) = &s {
BlockMode::from_str(s)
} else {
Err("Missing requires blockmode parameter".into())
}
}
fn decode_image(image: String) -> Result<(Vec<u8>, Box<dyn ImageInfo>), String> {
match image.split_at(image.find('.').unwrap_or(0)).1 {
".png" => {
let decoder = png::Decoder::new(
File::open(&image).map_err(|err| format!("Failed to read `{image}`: {err}"))?,
);
let mut reader = decoder
.read_info()
.map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?;
let mut result = Vec::with_capacity(reader.output_buffer_size());
result.resize(reader.output_buffer_size(), 0);
let info = reader
.next_frame(result.as_mut_slice())
.map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?;
result.resize(info.buffer_size(), 0);
Ok((result, Box::new(info)))
}
_ => Err(format!("Unable get image type for {image}")),
}
}
fn derive_seed(seed: &str) -> Result<[u8; 32], String> {
let mut result = [0u8; 32];
argon2::Argon2::default().hash_password_into(seed.as_bytes(), b"SEED SALT", &mut result)
.map_err(|err| format!("Failed to derive seed `{seed}`: {err}"))?;
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();
let mut data_pos = 0;
while read_data.len() < 9*8
{
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
{
data_pos = algorithm.read_block(&data, data_pos, &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
{
data_pos = algorithm.read_block(&data, data_pos, &mut read_data, &blockmode);
}
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);
}
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;
let embed_file = matches
.opt_str("i")
.ok_or(format!("Embed file is required"))?;
let (mut 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(),
)?;
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,
comment: matches.opt_str("c"),
};
let header_data = header.to_data(1, embed_data.len() as u32);
// Check length
if embed_data.len() + header_data.len() > max_size {
Err(format!(
"Cannot embed {}bytes into {}bytes using the {algorithm} algorithm with blockmode {}. Max embeddable size: {max_size}bytes",
embed_data.len()+header_data.len(),
data.len(),
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;
// 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;
let mut data_pos = 0;
for i in 0 .. blocks_num
{
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);
}
println!("{}", bv.len());
for i in 10..80 {
let b = (data[i*2] & 0b1111) | ((data[i*2+1] & 0b1111) << 4);
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>);
info.encode(w, data);
Ok(vec![])
}
*/
fn main() -> ExitCode {
let args: Vec<String> = env::args().collect();
let program = args[0].clone();
let mut opts = Options::new();
opts.optopt("i", "input", "Input file", "PATH");
opts.optflag("e", "encode", "Encode file");
opts.optflag("d", "decode", "Decode mode");
opts.optopt("c", "comment", "Header comment", "TXT");
opts.optopt("s", "seed", "Force seed", "TXT");
opts.optflag("", "no-crc", "Disables CRC");
opts.optflag("z", "info", "Read information");
opts.optopt("l", "algorithm", "Embed algorithm", "lo3");
opts.optflag("h", "help", "Print this help menu");
opts.optflag("v", "version", "Print program version and licenses");
let matches = match opts.parse(&args[1..]) {
Ok(m) => m,
Err(f) => {
panic!("{}", f.to_string())
}
};
if matches.opt_present("v") {
print_version();
return ExitCode::SUCCESS;
}
if matches.opt_present("h") {
print_usage(&program, opts);
return ExitCode::SUCCESS;
}
if matches.free.is_empty() {
print_usage(&program, opts);
return ExitCode::FAILURE;
}
/*
let input = matches.free[0].clone();
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) => {
eprintln!("{e}");
return ExitCode::FAILURE;
}
}
}
*/
ExitCode::SUCCESS
}

195
src/block.rs → src/png_embed/block.rs
View file

@ -1,6 +1,3 @@
use std::fmt::Formatter;
use std::str::FromStr;
use bitvec::slice::BitSlice; use bitvec::slice::BitSlice;
use bitvec::vec::BitVec; use bitvec::vec::BitVec;
use rand::prelude::SliceRandom; use rand::prelude::SliceRandom;
@ -8,80 +5,29 @@ use rand::Rng;
use crate::embed::EmbedAlgorithm; use crate::embed::EmbedAlgorithm;
/// Block mode for embedded data /// Gets the best blocksize (i.e. that minimize remaining space) for a certain data length.
#[derive(Debug)] /// The blocksize is a number in range [16, 65536]
pub struct BlockMode { pub fn best_blocksize(len: usize) -> usize {
pub len: usize, let mut best_remainder = len;
pub crc_len: usize, let mut best_p = 0;
} for p in 4..16 {
let remainder = len % (1 << p);
impl BlockMode { if remainder <= best_remainder {
/// Gets the best [`BlockMode`] and the remainder best_remainder = remainder;
pub fn from_length(len: usize, crc: bool) -> Self { best_p = p;
let mut best_remainder = len;
let mut best_p = 0;
for p in 4..16 {
let remainder = len % (1 << p);
if remainder <= best_remainder {
best_remainder = remainder;
best_p = p;
}
}
Self {
len: 1 << best_p,
crc_len: (best_p / 4) * crc as usize,
} }
} }
pub fn to_data(&self) -> u8 { 1 << best_p
((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)
}
}
impl FromStr for BlockMode {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let size = s
.parse::<usize>()
.map_err(|err| format!("Failed to parse `{}` as block size: {err}", s))?;
if size < 6 || size > 16 {
Err(format!(
"Invalid block size specified: `{size}` expected value within [6; 16]"
))?;
}
Ok(BlockMode {
len: 1 << size,
crc_len: size,
})
}
} }
/// Struct to hold the positions of data blocks
#[derive(Debug)] #[derive(Debug)]
pub struct BlockPlacement<'a> { pub struct BlockPlacement<'a> {
algorithm: &'a EmbedAlgorithm, algorithm: &'a EmbedAlgorithm,
data: &'a mut [u8], data: &'a mut [u8],
block_size: usize, block_size: usize,
blocks: Vec<usize>, pub blocks: Vec<usize>,
} }
impl<'a> BlockPlacement<'a> { impl<'a> BlockPlacement<'a> {
@ -91,8 +37,8 @@ impl<'a> BlockPlacement<'a> {
// //
// Will fail if the data is too small to hold all the blocks // Will fail if the data is too small to hold all the blocks
pub fn new<R>( pub fn new<R>(
data: &'a mut [u8],
algorithm: &'a EmbedAlgorithm, algorithm: &'a EmbedAlgorithm,
data: &'a mut [u8],
block_size: usize, block_size: usize,
embed_size: usize, embed_size: usize,
rng: &mut R, rng: &mut R,
@ -132,7 +78,35 @@ impl<'a> BlockPlacement<'a> {
}) })
} }
pub fn write_embed(&mut self) // Embeds the data into the original image
pub fn write_embed(&mut self, embed: &BitSlice<u8>) {
assert_eq!(embed.len() % 8, 0);
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 mut index = 0;
match self.algorithm {
EmbedAlgorithm::Lo(bits) => {
for block in &self.blocks {
for i in 0..self.block_size {
let pos = block * self.block_size + i;
let hi = std::cmp::min(*bits as usize, embed.len() - index);
self.data[pos] &= !((1 << hi) - 1);
self.data[pos] |= bits_to_byte(&embed[index..], hi as u8);
index += hi;
}
}
}
}
}
} }
// Iterator over blocks in the resulting image // Iterator over blocks in the resulting image
@ -145,9 +119,12 @@ pub struct BlockPlacementIterator<'a> {
index: usize, index: usize,
// Position of the blocks // Position of the blocks
blocks: Vec<usize>, blocks: Vec<usize>,
// Iterator over the current block
block_it: Option<BlockIterator<'a>>,
} }
impl<'a> BlockPlacementIterator<'a> { impl<'a> BlockPlacementIterator<'a> {
/// Creates a new embed iterator
pub fn new<R: Rng + ?Sized>( pub fn new<R: Rng + ?Sized>(
algorithm: &'a EmbedAlgorithm, algorithm: &'a EmbedAlgorithm,
data: &'a [u8], data: &'a [u8],
@ -163,29 +140,50 @@ impl<'a> BlockPlacementIterator<'a> {
// Shuffle the block order // Shuffle the block order
blocks.shuffle(rng); blocks.shuffle(rng);
let first_block_pos = blocks[0] * block_size;
let first_block = &data[first_block_pos..first_block_pos + block_size];
Self { Self {
algorithm, algorithm,
data, data,
block_size, block_size,
index: 0, index: 0,
blocks, blocks,
block_it: Some(BlockIterator::new(Block(algorithm, first_block), None)),
} }
} }
} }
impl<'a> Iterator for BlockPlacementIterator<'a> { impl<'a> Iterator for BlockPlacementIterator<'a> {
type Item = Block<'a>; type Item = u8;
/// Gets the next embedded byte in the image
///
/// # Note
///
/// Even when the [`next()`] is Some(..), if the iterator is past the embed's length, it will
/// return garbage data.
fn next(&mut self) -> Option<Self::Item> { fn next(&mut self) -> Option<Self::Item> {
if self.index == self.blocks.len() { self.block_it.as_ref()?;
return None;
if let Some(byte) = self.block_it.as_mut().unwrap().next() {
Some(byte)
} else {
self.index += 1;
// Get next block
if self.index == self.blocks.len() {
return None;
}
let block_pos = self.blocks[self.index] * self.block_size;
let block = &self.data[block_pos..block_pos + self.block_size];
self.block_it = Some(BlockIterator::new(
Block(self.algorithm, block),
self.block_it.take(),
));
self.next()
} }
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))
} }
} }
@ -196,7 +194,7 @@ pub struct Block<'a>(&'a EmbedAlgorithm, &'a [u8]);
// Iterator to read embedded data inside a block // Iterator to read embedded data inside a block
pub struct BlockIterator<'a> { pub struct BlockIterator<'a> {
// Block of the iterator // Block of the iterator
block: &'a Block<'a>, block: Block<'a>,
// Byte position in [`data`] // Byte position in [`data`]
index: usize, index: usize,
@ -206,7 +204,7 @@ pub struct BlockIterator<'a> {
} }
impl<'a> BlockIterator<'a> { impl<'a> BlockIterator<'a> {
pub fn new(block: &'a Block, previous: Option<BlockIterator>) -> Self { pub fn new(block: Block<'a>, previous: Option<BlockIterator>) -> Self {
if let Some(previous) = previous { if let Some(previous) = previous {
Self { Self {
block, block,
@ -275,6 +273,28 @@ mod tests {
use super::*; use super::*;
#[test]
fn test_write() {
let algorithm = EmbedAlgorithm::Lo(2);
let mut data = vec![0u8; 8];
let embed = vec![0xFF, 0xFF];
let embed_bits = BitVec::<u8>::from_slice(embed.as_slice());
let mut rand = ChaCha8Rng::from_seed([1u8; 32]);
let mut placement = BlockPlacement::new::<_>(
&algorithm,
data.as_mut_slice(),
4,
embed_bits.len() / 8,
&mut rand,
)
.unwrap();
placement.write_embed(embed_bits.as_bitslice());
assert_eq!(data, vec![0b00000011; 8]);
}
#[test] #[test]
fn block_iterator() { fn block_iterator() {
let algorithm = EmbedAlgorithm::Lo(3); let algorithm = EmbedAlgorithm::Lo(3);
@ -284,17 +304,17 @@ mod tests {
]; ];
let block = Block(&algorithm, &data); let block = Block(&algorithm, &data);
let mut it = BlockIterator::new(&block, None); let mut it = BlockIterator::new(block, None);
assert_eq!(it.next(), Some(0b10_001_000)); assert_eq!(it.next(), Some(0b10_001_000));
assert_eq!(it.next(), Some(0b0_100_011_0)); assert_eq!(it.next(), Some(0b0100_0110));
assert_eq!(it.next(), Some(0b000_111_11)); assert_eq!(it.next(), Some(0b0001_1111));
} }
#[test] #[test]
fn blockplacement_iterator() { fn blockplacement_iterator() {
let algorithm = EmbedAlgorithm::Lo(4); let algorithm = EmbedAlgorithm::Lo(4);
let mut data = vec![ let data = vec![
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 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, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
0x1c, 0x1d, 0x1e, 0x1f, 0x1c, 0x1d, 0x1e, 0x1f,
@ -307,9 +327,10 @@ mod tests {
let mut positions = (0..8).collect::<Vec<_>>(); let mut positions = (0..8).collect::<Vec<_>>();
positions.shuffle(&mut rand); positions.shuffle(&mut rand);
for i in 0..8 { for i in 0..data.len() / 2 {
let block = it.next().unwrap(); let byte = it.next().unwrap();
assert_eq!(block.1[0] / 4, positions[i]); // TODO...
//assert_eq!(byte, data[positions[i/4]*4+(i%4)]);
} }
} }
} }

54
src/png_embed/embed.rs Normal file
View file

@ -0,0 +1,54 @@
use std::fmt::Formatter;
use std::str::FromStr;
/// Algorithm to embed data
#[derive(Debug)]
pub enum EmbedAlgorithm {
Lo(u8),
}
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,
}
}
}
impl core::fmt::Display for EmbedAlgorithm {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
EmbedAlgorithm::Lo(bits) => write!(f, "Lo({bits})"),
}
}
}
impl FromStr for EmbedAlgorithm {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let (dig_pos, _) = s
.char_indices()
.find(|(_, c)| c.is_ascii_digit())
.ok_or(format!("Unknown algorithm: {s}"))?;
let (first, second) = s.split_at(dig_pos);
match first {
"lo" => {
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 > 7 || value == 0 {
Err(format!(
"Cannot specify {value} bits for `lo` method, must be within [1, 7]"
))
} else {
Ok(EmbedAlgorithm::Lo(value))
}
}
_ => Err(format!("Unknown algorithm: {s}")),
}
}
}

139
src/png_embed/header.rs Normal file
View file

@ -0,0 +1,139 @@
use crc::Crc;
use crate::block::BlockPlacementIterator;
/// The program's version.
/// Used for compatibility reasons.
#[repr(u16)]
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy)]
pub enum Version {
VERSION_1,
}
impl TryFrom<u16> for Version {
type Error = String;
fn try_from(value: u16) -> Result<Self, Self::Error> {
match value {
0 => Ok(Version::VERSION_1),
ver => Err(format!("Unknown version: {ver}")),
}
}
}
pub trait Encode {
/// Encode the data into a vector
fn encode(&self, vec: &mut Vec<u8>);
}
pub trait Decode {
type Type;
/// Decode the data from an iterator
fn decode(it: &mut BlockPlacementIterator) -> Result<Self::Type, String>;
}
/// Embed data header:
/// +---------+----------+----------+-------------+---------+
/// | Version | Data Len | Data CRC | Comment Len | Comment |
/// +---------+----------+----------+-------------+---------+
/// | 2 | 4 | 4 | 2 | varies |
/// +---------+----------+----------+-------------+---------+
#[derive(Debug)]
pub struct Header {
pub version: Version,
pub data_len: u32,
pub data_crc: u32,
pub comment: Option<String>,
}
impl Header {
/// Construct a new header from the embedded data
pub fn new(version: Version, data: &[u8], comment: Option<String>) -> Result<Self, String> {
if data.len() > u32::MAX as usize {
return Err(format!(
"Embedded data length: {} is greater than maximum {}",
data.len(),
u32::MAX
));
} else if let Some(len) = comment.as_ref().map(|c| c.len()) {
if len > u16::MAX as usize {
return Err(format!(
"Embedded comment is too long, maximum length: {}, got {len}",
u16::MAX
));
}
}
Ok(Self {
version,
data_len: data.len() as u32,
data_crc: Crc::<u32>::new(&crc::CRC_32_CKSUM).checksum(data),
comment,
})
}
}
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());
// 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());
}
}
}
impl Decode for Header {
type Type = Header;
fn decode(it: &mut BlockPlacementIterator) -> Result<Header, String> {
let mut count = 0;
let mut next = || -> Result<u8, String> {
let result = it
.next()
.ok_or(format!("Failed to get byte at index: {count}"));
count += 1;
result
};
let version = u16::from_le_bytes([next()?, next()?]);
let data_len = u32::from_le_bytes([next()?, next()?, next()?, next()?]);
let data_crc = u32::from_le_bytes([next()?, next()?, next()?, next()?]);
let comment_length = u16::from_le_bytes([next()?, next()?]);
let comment = if comment_length != 0 {
let mut comment_data = Vec::with_capacity(comment_length as usize);
for _ in 0..comment_length {
comment_data.push(next()?);
}
Some(
String::from_utf8(comment_data)
.map_err(|e| format!("Failed to retrieve comment: {e}"))?,
)
} else {
None
};
Ok(Header {
version: Version::try_from(version)?,
data_len,
data_crc,
comment,
})
}
}

0
src/image.rs → src/png_embed/image.rs
View file

335
src/png_embed/main.rs Normal file
View file

@ -0,0 +1,335 @@
pub mod block;
pub mod embed;
pub mod header;
pub mod image;
use std::env;
use std::fs::File;
use std::io::BufWriter;
use std::io::Write;
use std::process::ExitCode;
use std::str::FromStr;
use bitvec::prelude::*;
use block::best_blocksize;
use block::BlockPlacement;
use block::BlockPlacementIterator;
use crc::Crc;
use embed::EmbedAlgorithm;
use getopts::Matches;
use getopts::Options;
use header::Decode;
use header::Encode;
use header::Header;
use image::ImageInfo;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
fn print_usage(program: &str, opts: Options) {
let brief = format!(
"Usage: {0} -l ALGORITHM -(e|d|z) [EMBED] FILE -o OUTPUT [opts]
Encode: {0} -l lo3 -e embed.jpg input.png -o out.png -c \"Embedded JPEG file\"
Info: {0} -l lo3 out.png # Embedded JPEG file
Decode: {0} -l lo3 out.png > decoded.jpg",
program
);
print!("{}", opts.usage(&brief));
}
fn print_version() {
print!(
r#"png_embed (c) ef3d0c3e -- Embed data into PNGs
Copyright (c) 2024
NML is licensed under the GNU Affero General Public License version 3 (AGPLv3),
under the terms of the Free Software Foundation <https://www.gnu.org/licenses/agpl-3.0.en.html>.
This program is free software; you may modify and redistribute it.
There is NO WARRANTY, to the extent permitted by law."#
);
}
impl ImageInfo for png::OutputInfo {
fn width(&self) -> u32 { self.width }
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();
}
}
fn decode_image(image: &str) -> Result<(Vec<u8>, Box<dyn ImageInfo>), String> {
match image.split_at(image.find('.').unwrap_or(0)).1 {
".png" => {
let decoder = png::Decoder::new(
File::open(image).map_err(|err| format!("Failed to read `{image}`: {err}"))?,
);
let mut reader = decoder
.read_info()
.map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?;
let mut result = vec![0; reader.output_buffer_size()];
let info = reader
.next_frame(result.as_mut_slice())
.map_err(|err| format!("Failed to read png info for `{image}`: {err}"))?;
result.resize(info.buffer_size(), 0);
Ok((result, Box::new(info)))
}
_ => Err(format!("Unable get image type for {image}")),
}
}
// Derives the seed from a given string.
// Currently using Argon with salt: `png_data embed`
fn derive_seed(seed: &str) -> Result<[u8; 32], String> {
let mut result = [0u8; 32];
argon2::Argon2::default()
.hash_password_into(seed.as_bytes(), b"png_data embed", &mut result)
.map_err(|err| format!("Failed to derive seed `{seed}`: {err}"))?;
Ok(result)
}
fn encode(
input: String,
embed: String,
output: String,
algorithm: String,
matches: Matches,
) -> Result<(), String> {
let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?;
let (mut data, info) = decode_image(input.as_str())?;
let block_size = best_blocksize(info.size());
let seed = derive_seed(
matches
.opt_str("s")
.unwrap_or(format!("{}x{}", info.width(), info.height()))
.as_str(),
)?;
let comment = matches.opt_str("c");
// Data
let embed_file_data = std::fs::read(&embed)
.map_err(|err| format!("Failed to read embed file `{embed}`: {err}"))?;
// Header
let header = Header::new(
header::Version::VERSION_1,
embed_file_data.as_slice(),
comment,
)?;
// Result
let mut embed_data = vec![];
header.encode(&mut embed_data);
embed_data.extend(embed_file_data);
eprintln!("=== HEADER ===");
eprintln!("Version: {:#?}", header.version);
eprintln!(
"Comment: {}",
header.comment.as_ref().map_or("", |c| c.as_str())
);
eprintln!("Data: {}bytes CRC[{:X}]", header.data_len, header.data_crc);
eprintln!("Block: {block_size}bytes");
let mut rand = ChaCha8Rng::from_seed(seed);
let mut placement = BlockPlacement::new(
&algorithm,
data.as_mut_slice(),
block_size,
embed_data.len(),
&mut rand,
)?;
eprintln!("Required blocks: {}", placement.blocks.len());
eprintln!("==============");
placement.write_embed(embed_data.as_slice().view_bits::<Lsb0>());
let outfile = File::create(&output).unwrap();
let w = &mut BufWriter::new(Box::new(outfile) as Box<dyn Write>);
info.encode(w, data);
Ok(())
}
fn decode_header(input: String, algorithm: String, matches: Matches) -> Result<(), String> {
let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?;
let (data, info) = decode_image(input.as_str())?;
let block_size = best_blocksize(info.size());
let seed = derive_seed(
matches
.opt_str("s")
.unwrap_or(format!("{}x{}", info.width(), info.height()))
.as_str(),
)?;
let mut rand = ChaCha8Rng::from_seed(seed);
let mut it = BlockPlacementIterator::new(&algorithm, data.as_slice(), block_size, &mut rand);
let header = Header::decode(&mut it)?;
eprintln!("=== HEADER ===");
eprintln!("Version: {:#?}", header.version);
eprintln!(
"Comment: \"{}\"",
header.comment.as_ref().map_or("", |c| c.as_str())
);
eprintln!("Data: {}bytes CRC[{:X}]", header.data_len, header.data_crc);
eprintln!("==============");
Ok(())
}
fn decode(
input: String,
output: String,
algorithm: String,
matches: Matches,
) -> Result<(), String> {
let algorithm = EmbedAlgorithm::from_str(algorithm.as_str())?;
let (data, info) = decode_image(input.as_str())?;
let block_size = best_blocksize(info.size());
let seed = derive_seed(
matches
.opt_str("s")
.unwrap_or(format!("{}x{}", info.width(), info.height()))
.as_str(),
)?;
let mut rand = ChaCha8Rng::from_seed(seed);
let mut it = BlockPlacementIterator::new(&algorithm, data.as_slice(), block_size, &mut rand);
let header = Header::decode(&mut it)?;
let mut data = Vec::with_capacity(header.data_len as usize);
while data.len() < header.data_len as usize {
data.push(
it.next()
.ok_or(format!("Failed to read data byte at {}", data.len()))?,
);
}
// Check CRC
let data_crc = Crc::<u32>::new(&crc::CRC_32_CKSUM).checksum(data.as_slice());
if data_crc != header.data_crc {
Err(format!(
"Data CRC do not match: HEADER={:X} GOT={data_crc:X}",
header.data_crc
))?;
}
let outfile = File::create(&output)
.map_err(|e| format!("Failed to create output file `{output}`: {e}"))?;
let w = &mut BufWriter::new(Box::new(outfile) as Box<dyn Write>);
w.write_all(data.as_slice())
.map_err(|e| format!("Failed to write to output file `{output}`: {e}"))?;
eprintln!("File written to `{output}`");
Ok(())
}
fn main() -> ExitCode {
let args: Vec<String> = env::args().collect();
let program = args[0].clone();
let mut opts = Options::new();
opts.optopt("e", "embed", "Embed file", "PATH");
opts.optopt("o", "output", "Output file", "PATH");
opts.optflag("d", "decode", "Decode mode");
opts.optopt("c", "comment", "Header comment", "TXT");
opts.optopt(
"s",
"seed",
"Force a seed, defaults to \"{width}x{height}\"",
"TXT",
);
opts.optflag("z", "info", "Read header");
opts.optopt("l", "algorithm", "Embed algorithm", "lo3");
opts.optflag("h", "help", "Print this help menu");
opts.optflag("v", "version", "Print program version and licenses");
let matches = match opts.parse(&args[1..]) {
Ok(m) => m,
Err(f) => {
panic!("{}", f.to_string())
}
};
if matches.opt_present("v") {
print_version();
return ExitCode::SUCCESS;
}
if matches.opt_present("h") {
print_usage(&program, opts);
return ExitCode::SUCCESS;
}
// Get input file
if matches.free.is_empty() {
eprintln!("Missing input file");
print_usage(&program, opts);
return ExitCode::FAILURE;
}
let input_file = matches.free[0].clone();
// Check options
if matches.opt_present("e") as usize
+ matches.opt_present("d") as usize
+ matches.opt_present("z") as usize
> 1
{
eprintln!("Specify either `-e(--embed)`, -z(--info) or `-d(--decode)`");
return ExitCode::FAILURE;
} else if !matches.opt_present("l") {
eprintln!("Missing algorithm name");
return ExitCode::FAILURE;
}
// Get algorithm
let algorithm = matches.opt_str("l").unwrap();
if matches.opt_present("e") {
let embed_file = matches.opt_str("e").unwrap();
if !matches.opt_present("o") {
eprintln!("Missing -o(utput) file");
return ExitCode::FAILURE;
}
let output_file = matches.opt_str("o").unwrap();
if let Err(e) = encode(input_file, embed_file, output_file, algorithm, matches) {
eprintln!("{e}");
return ExitCode::FAILURE;
}
} else if matches.opt_present("z") {
if let Err(e) = decode_header(input_file, algorithm, matches) {
eprintln!("{e}");
return ExitCode::FAILURE;
}
} else if matches.opt_present("d") {
if !matches.opt_present("o") {
eprintln!("Missing -o(utput) file");
return ExitCode::FAILURE;
}
let output_file = matches.opt_str("o").unwrap();
if let Err(e) = decode(input_file, output_file, algorithm, matches) {
eprintln!("{e}");
return ExitCode::FAILURE;
}
} else {
print_usage(&program, opts);
return ExitCode::FAILURE;
}
ExitCode::SUCCESS
}