pub mod error;
mod frame;
use std::path::Path;
use zeroize::Zeroizing;
use crate::cost::hill::{CostError, HillCostProvider};
use crate::cost::CostProvider;
use crate::crypto::aead::{
compress_and_encrypt, decrypt_and_decompress, AEADCipher, AEADError, CryptoError,
XChaCha20Poly1305Cipher,
};
use crate::crypto::expand::expand_master_key;
use crate::crypto::kdf::{Argon2Kdf, KeyDeriver};
use crate::image_io::buffer::{CoverSource, ImageBuffer};
use crate::image_io::phash::{
compute_stable_phash, phash_salt_hypotheses, recover_phash_salt, PHashError, PHashSalt,
};
use crate::image_io::validate::load_and_validate;
use crate::stego::permute::generate_pixel_permutation;
use crate::stego::sizer::{compute_capacity, validate_payload_fits, EmbeddingMode, SizerError};
use crate::stego::stc::{stc_decode_safe, stc_encode_safe, StcConfig};
pub use crate::pipeline::error::{OutputError, PipelineError};
const PROVISIONAL_PREFIX_BYTES: usize = 64;
#[derive(Debug)]
pub struct EmbedReport {
pub pixels_modified: usize,
pub payload_bytes: usize,
pub effective_bpp: f32,
pub image_dimensions: (u32, u32),
}
#[derive(Debug)]
pub struct ExtractReport {
pub payload_bytes: usize,
}
pub struct EmbedPipeline<KDF, AEAD, COST> {
kdf: KDF,
aead: AEAD,
cost: COST,
}
impl<KDF, AEAD, COST> EmbedPipeline<KDF, AEAD, COST> {
pub fn new(kdf: KDF, aead: AEAD, cost: COST) -> Self {
Self { kdf, aead, cost }
}
}
impl EmbedPipeline<Argon2Kdf, XChaCha20Poly1305Cipher, HillCostProvider> {
pub fn default_secure() -> Self {
Self::new(
Argon2Kdf::default_secure(),
XChaCha20Poly1305Cipher::new(),
HillCostProvider::new(),
)
}
}
enum Attempt {
Recovered {
plaintext: Zeroizing<Vec<u8>>,
ciphertext_bytes: usize,
},
Rejected(Zeroizing<Vec<u8>>),
}
impl<KDF, AEAD, COST> EmbedPipeline<KDF, AEAD, COST>
where
KDF: KeyDeriver,
AEAD: AEADCipher,
COST: CostProvider<Error = CostError>,
{
pub fn embed(
&self,
image_path: &Path,
plaintext: Zeroizing<Vec<u8>>,
password: Zeroizing<Vec<u8>>,
output_path: &Path,
) -> Result<EmbedReport, PipelineError> {
let mut image_buffer = load_and_validate(image_path)?;
let image_dimensions = image_buffer.dimensions();
let pixel_count = image_buffer.pixel_count();
let phash_salt = compute_stable_phash(&image_buffer)?;
let master_key = self.kdf.derive(password.as_slice(), &phash_salt)?;
drop(password);
drop(phash_salt);
let derived_keys = expand_master_key(&master_key)?;
drop(master_key);
let ciphertext = compress_and_encrypt(
plaintext.as_slice(),
derived_keys.enc_key(),
derived_keys.nonce(),
&self.aead,
)?;
drop(plaintext);
let cost_map = self.cost.compute(&image_buffer)?;
let capacity = compute_capacity(&cost_map, EmbeddingMode::Symmetric);
validate_payload_fits(ciphertext.len() + frame::FRAME_OVERHEAD_BYTES, &capacity)?;
let permutation = generate_pixel_permutation(pixel_count, derived_keys.stc_seed());
let mut cover_symbols = frame::gather_cover_symbols(&image_buffer, &permutation);
let cost_reordered = frame::reorder_costs(cost_map.costs(), &permutation);
drop(cost_map);
let length_header = frame::encode_length_header(ciphertext.len()).ok_or_else(|| {
SizerError::PayloadTooLarge {
payload: ciphertext.len(),
available: u32::MAX as usize,
deficit: ciphertext.len().saturating_sub(u32::MAX as usize),
}
})?;
let stc_config = StcConfig::new(*derived_keys.stc_seed());
let (header_costs, payload_costs) = frame::split_regions(&cost_reordered);
let (header_cover, payload_cover) = frame::split_regions_mut(&mut cover_symbols);
let header_changes =
stc_encode_safe(header_cover, header_costs, &length_header, &stc_config)?;
let payload_changes = stc_encode_safe(
payload_cover,
payload_costs,
ciphertext.as_slice(),
&stc_config,
)?;
let payload_bytes = ciphertext.len();
drop(ciphertext);
drop(derived_keys);
drop(stc_config);
drop(cost_reordered);
frame::apply_cover_symbols(&mut image_buffer, &permutation, &cover_symbols);
drop(permutation);
drop(cover_symbols);
frame::write_png(&image_buffer, output_path)?;
let embedded_bits = frame::LENGTH_HEADER_BITS + payload_bytes * 8;
Ok(EmbedReport {
pixels_modified: header_changes + payload_changes,
payload_bytes,
effective_bpp: embedded_bits as f32 / pixel_count.max(1) as f32,
image_dimensions,
})
}
pub fn extract(
&self,
stego_path: &Path,
password: Zeroizing<Vec<u8>>,
) -> Result<(Zeroizing<Vec<u8>>, ExtractReport), PipelineError> {
let stego_image = load_and_validate(stego_path)?;
let hypotheses = phash_salt_hypotheses(&stego_image)?;
match self.attempt_extract(&stego_image, &hypotheses.primary, password.as_slice())? {
Attempt::Recovered {
plaintext,
ciphertext_bytes,
} => {
drop(password);
Ok((
plaintext,
ExtractReport {
payload_bytes: ciphertext_bytes,
},
))
}
Attempt::Rejected(prefix) => {
let Some(alternative) = hypotheses.alternative else {
drop(password);
return Err(PipelineError::Crypto(CryptoError::AEADError(
AEADError::AuthenticationFailed,
)));
};
let verdict = recover_phash_salt(
&stego_image,
password.as_slice(),
&self.kdf,
prefix.as_slice(),
);
drop(prefix);
let outcome = match verdict {
Ok(confirmed) => {
drop(confirmed);
drop(password);
return Err(PipelineError::Crypto(CryptoError::AEADError(
AEADError::AuthenticationFailed,
)));
}
Err(PHashError::RecoveryFailed) => {
self.attempt_extract(&stego_image, &alternative, password.as_slice())
}
Err(err) => Err(PipelineError::PHash(err)),
};
drop(password);
match outcome? {
Attempt::Recovered {
plaintext,
ciphertext_bytes,
} => Ok((
plaintext,
ExtractReport {
payload_bytes: ciphertext_bytes,
},
)),
Attempt::Rejected(_) => Err(PipelineError::Crypto(CryptoError::AEADError(
AEADError::AuthenticationFailed,
))),
}
}
}
}
fn attempt_extract(
&self,
stego_image: &ImageBuffer,
salt: &PHashSalt,
password: &[u8],
) -> Result<Attempt, PipelineError> {
let master_key = self.kdf.derive(password, salt)?;
let derived_keys = expand_master_key(&master_key)?;
drop(master_key);
let permutation =
generate_pixel_permutation(stego_image.pixel_count(), derived_keys.stc_seed());
let cover_symbols = frame::gather_cover_symbols(stego_image, &permutation);
drop(permutation);
let stc_config = StcConfig::new(*derived_keys.stc_seed());
let (header_region, payload_region) = frame::split_regions(&cover_symbols);
let header = stc_decode_safe(header_region, frame::LENGTH_HEADER_BITS, &stc_config)?;
let announced = frame::decode_length_header(&header).unwrap_or(0);
if announced < frame::MIN_CIPHERTEXT_BYTES
|| announced.saturating_mul(8) > stc_config.capacity_bits(payload_region.len())
{
return Ok(Attempt::Rejected(Zeroizing::new(Vec::new())));
}
let ciphertext =
Zeroizing::new(stc_decode_safe(payload_region, announced * 8, &stc_config)?);
drop(cover_symbols);
drop(stc_config);
let outcome = decrypt_and_decompress(
ciphertext.as_slice(),
derived_keys.enc_key(),
derived_keys.nonce(),
&self.aead,
);
drop(derived_keys);
match outcome {
Ok(plaintext) => Ok(Attempt::Recovered {
plaintext,
ciphertext_bytes: announced,
}),
Err(CryptoError::AEADError(_)) => Ok(Attempt::Rejected(Zeroizing::new(
ciphertext
.iter()
.copied()
.take(PROVISIONAL_PREFIX_BYTES)
.collect(),
))),
Err(err) => Err(PipelineError::Crypto(err)),
}
}
}