png_data working
This commit is contained in:
parent
2601e7dbb1
commit
892f22ce1d
3 changed files with 229 additions and 11 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ pub trait Decode {
|
|||
|
||||
/// Decode the data from an iterator
|
||||
fn decode<I>(it: &mut I) -> Result<Self::Type, String>
|
||||
where I: Iterator<Item = u8>;
|
||||
where I: Iterator<Item = (usize, u8)>;
|
||||
}
|
||||
|
||||
/// The program's version.
|
||||
|
|
@ -95,14 +95,14 @@ impl Decode for Header {
|
|||
type Type = Header;
|
||||
|
||||
fn decode<I>(it: &mut I) -> Result<Self::Type, String>
|
||||
where I: Iterator<Item = u8> {
|
||||
where I: Iterator<Item = (usize, u8)> {
|
||||
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
|
||||
result.map(|(_, b)| b)
|
||||
};
|
||||
|
||||
let version = u16::from_le_bytes([next()?, next()?]);
|
||||
|
|
|
|||
|
|
@ -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::<u8>())
|
||||
}
|
||||
|
||||
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::<u32>::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
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue