diff --git a/Cargo.toml b/Cargo.toml index e36c9ad..ba5b4ab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,10 @@ edition = "2021" name = "png_embed" path = "src/png_embed/main.rs" +[[bin]] +name = "png_data" +path = "src/png_data/main.rs" + [dependencies] argon2 = "0.5.3" bitvec = "1.0.1" diff --git a/src/png_data/header.rs b/src/png_data/header.rs index aafd546..b487ad6 100644 --- a/src/png_data/header.rs +++ b/src/png_data/header.rs @@ -11,7 +11,7 @@ pub trait Decode { /// Decode the data from an iterator fn decode(it: &mut I) -> Result - where I: Iterator; + where I: Iterator; } /// The program's version. @@ -95,14 +95,14 @@ impl Decode for Header { type Type = Header; fn decode(it: &mut I) -> Result - where I: Iterator { + where I: Iterator { let mut count = 0; let mut next = || -> Result { let result = it .next() .ok_or(format!("Failed to get byte at index: {count}")); count += 1; - result + result.map(|(_, b)| b) }; let version = u16::from_le_bytes([next()?, next()?]); diff --git a/src/png_data/main.rs b/src/png_data/main.rs index 5c796e5..5496f5a 100644 --- a/src/png_data/main.rs +++ b/src/png_data/main.rs @@ -1,10 +1,22 @@ +#![feature(isqrt)] mod header; use std::env; +use std::fs::File; +use std::io::BufWriter; use std::process::ExitCode; +use crc::Crc; use getopts::Matches; use getopts::Options; +use header::Decode; +use header::Encode; +use header::Header; +use png::BitDepth; +use png::ColorType; +use rand::Rng; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; fn print_usage(program: &str, opts: Options) { let brief = format!( @@ -29,13 +41,195 @@ There is NO WARRANTY, to the extent permitted by law."# ); } -fn best_layout(size: u32, bits_per_pixel: u8) -> (u32, u32) { - let sz : f64 = size as f64 / bits_per_pixel as f64; - let width = sz.sqrt().floor(); - (width as u32, (sz / width as f64).ceil() as u32) +fn str_to_layout(layout: &str) -> Result<(ColorType, BitDepth), String> { + let split = layout + .char_indices() + .find(|(_, c)| c.is_ascii_digit()) + .ok_or(format!("Unable to find number for layout's bit depth"))? + .0; + match layout.split_at(split) { + ("rgb", bits) => match bits { + "8" => Ok((ColorType::Rgb, BitDepth::Eight)), + "16" => Ok((ColorType::Rgb, BitDepth::Sixteen)), + _ => Err(format!("Color type rgb cannot have bit depth: {bits}")), + }, + ("rgba", bits) => match bits { + "8" => Ok((ColorType::Rgba, BitDepth::Eight)), + "16" => Ok((ColorType::Rgba, BitDepth::Sixteen)), + _ => Err(format!("Color type rgba cannot have bit depth: {bits}")), + }, + ("g", bits) => match bits { + "1" => Ok((ColorType::Grayscale, BitDepth::One)), + "2" => Ok((ColorType::Grayscale, BitDepth::Two)), + "4" => Ok((ColorType::Grayscale, BitDepth::Four)), + "8" => Ok((ColorType::Grayscale, BitDepth::Eight)), + "16" => Ok((ColorType::Grayscale, BitDepth::Sixteen)), + _ => Err(format!( + "Color type grayscale cannot have bit depth: {bits}" + )), + }, + ("ga", bits) => match bits { + "1" => Ok((ColorType::GrayscaleAlpha, BitDepth::One)), + "2" => Ok((ColorType::GrayscaleAlpha, BitDepth::Two)), + "4" => Ok((ColorType::GrayscaleAlpha, BitDepth::Four)), + "8" => Ok((ColorType::GrayscaleAlpha, BitDepth::Eight)), + "16" => Ok((ColorType::GrayscaleAlpha, BitDepth::Sixteen)), + _ => Err(format!( + "Color type grayscale alpha cannot have bit depth: {bits}" + )), + }, + _ => Err(format!("Uknown layout: {layout}")), + } +} + +fn bits_per_pixel(colors: ColorType, depth: BitDepth) -> u8 { + match colors { + ColorType::Rgb => depth as u8 * 3, + ColorType::Rgba => depth as u8 * 4, + ColorType::Grayscale => depth as u8, + ColorType::GrayscaleAlpha => depth as u8 * 2, + _ => panic!("Unsupported color type: {colors:#?}"), + } +} + +fn best_layout(size: u64, bits_per_pixel: u8) -> (u32, u32) { + let sz = (size * 8).div_ceil(bits_per_pixel as u64); + let width = sz.isqrt(); + (width as u32, sz.div_ceil(width) as u32) } fn encode(input: String, output: String, layout: String, matches: Matches) -> Result<(), String> { + let layout = str_to_layout(layout.as_str())?; + let comment = matches.opt_str("c"); + + // Input file data + let input_data = std::fs::read(&input) + .map_err(|err| format!("Failed to read input file `{input}`: {err}"))?; + + // Header + let header = Header::new(header::Version::VERSION_1, input_data.as_slice(), comment)?; + let mut data = vec![]; + header.encode(&mut 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!("=============="); + + let bits_per_pixel = bits_per_pixel(layout.0, layout.1); + let (width, height) = best_layout( + (data.len() + input_data.len()) as u64, + bits_per_pixel + ); + + // Encode + let output_file = File::create(&output) + .map_err(|err| format!("Failed to open output file `{output}`: {err}"))?; + let ref mut w = BufWriter::new(output_file); + let mut encoder = png::Encoder::new(w, width, height); + encoder.set_color(layout.0); + encoder.set_depth(layout.1); + encoder.set_compression(png::Compression::Best); + let mut writer = encoder + .write_header() + .map_err(|err| format!("Failed to write png header: {err}"))?; + + // Image byte length + let byte_len = ((width as usize) * (height as usize) * (bits_per_pixel as usize)).div_ceil(8); + data.reserve(byte_len); + + data.extend_from_slice(input_data.as_slice()); + + // Fill with random data + let mut rng = ChaCha8Rng::from_entropy(); + while data.len() < byte_len { + data.push(rng.gen::()) + } + + writer + .write_image_data(&data) + .map_err(|err| format!("Failed to write image data: {err}"))?; + println!("File written to `{output}`"); + + Ok(()) +} + +fn decode_header(input: String, _matches: Matches) -> Result<(), String> { + // Input file data + let decoder = png::Decoder::new( + File::open(&input).map_err(|err| format!("Failed to read input file `{input}`: {err}"))?, + ); + let mut reader = decoder + .read_info() + .map_err(|err| format!("Failed to read png info for `{input}`: {err}"))?; + let mut data = vec![0; reader.output_buffer_size()]; + let info = reader + .next_frame(data.as_mut_slice()) + .map_err(|err| format!("Failed to read png info for `{input}`: {err}"))?; + + data.resize(info.buffer_size(), 0); + + + let mut it = data.iter().enumerate().map(|(idx, byte)| (idx, *byte)); + let header = Header::decode(&mut it).map_err(|err| format!("Failed to decode header: {err}"))?; + 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, _matches: Matches) -> Result<(), String> { + // Input file data + let decoder = png::Decoder::new( + File::open(&input).map_err(|err| format!("Failed to read input file `{input}`: {err}"))?, + ); + let mut reader = decoder + .read_info() + .map_err(|err| format!("Failed to read png info for `{input}`: {err}"))?; + let mut data = vec![0; reader.output_buffer_size()]; + let info = reader + .next_frame(data.as_mut_slice()) + .map_err(|err| format!("Failed to read png info for `{input}`: {err}"))?; + + data.resize(info.buffer_size(), 0); + + + let mut it = data.iter().enumerate().map(|(idx, byte)| (idx, *byte)); + let header = + { + //let mut temp_it = std::mem::take(&mut it); + Header::decode(&mut it).map_err(|err| format!("Failed to decode header: {err}"))? + }; + 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!("=============="); + + // Check crc + let data_start = it.next().ok_or(format!("Failed to get data start"))?.0; + let file_data = &data[data_start..data_start+header.data_len as usize]; + let crc = Crc::::new(&crc::CRC_32_CKSUM).checksum(file_data); + if crc != header.data_crc { + Err(format!("Data CRC[{crc:X}] does not match header CRC[{:X}]", header.data_crc))?; + } + + std::fs::write(&output, file_data).map_err(|err| format!("Failed to write to output file `{output}`: {err}"))?; + println!("File written to `{output}`"); + Ok(()) } @@ -44,11 +238,10 @@ fn main() -> ExitCode { let program = args[0].clone(); let mut opts = Options::new(); - opts.optopt("e", "encode", "Embed file", "PATH"); - opts.optopt("d", "decode", "Decode mode", "PATH"); - opts.optflag("z", "info", "Read header"); + opts.optopt("e", "encode", "Embed file", "FILE"); + opts.optopt("d", "decode", "Decode mode", "FILE"); + opts.optopt("z", "info", "Read header", "FILE"); opts.optopt("l", "layout", "Png image layout", "TXT"); - opts.optopt("p", "password", "Data password", "TXT"); opts.optopt("o", "output", "Output file", "PATH"); opts.optopt("c", "comment", "Header comment", "TXT"); opts.optflag("h", "help", "Print this help menu"); @@ -100,6 +293,27 @@ fn main() -> ExitCode { eprintln!("{e}"); return ExitCode::FAILURE; } + } else if let Some(input_file) = matches.opt_str("z") { + if let Err(e) = decode_header(input_file, matches) { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } else if let Some(input_file) = matches.opt_str("d") { + let output_file = match matches.opt_str("o") { + None => { + eprintln!("Missing required output (-o|--output) option"); + return ExitCode::FAILURE; + } + Some(output_file) => output_file, + }; + + if let Err(e) = decode(input_file, output_file, matches) { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + } else { + print_usage(&program, opts); + return ExitCode::SUCCESS; } ExitCode::SUCCESS }