1use crate::codec::jpeg::JpegImage;
61use crate::stego::armor::ecc;
62use crate::stego::crypto::{self, NONCE_LEN, SALT_LEN};
63use crate::stego::error::StegoError;
64use crate::stego::frame;
65use crate::stego::payload::{self, FileEntry, PayloadData};
66use crate::stego::permute::CoeffPos;
67use super::pipeline::{flat_get, flat_set};
68use crate::stego::side_info::nsf5_modify_coefficient;
69
70const SHADOW_FRAME_OVERHEAD: usize = 2 + SALT_LEN + NONCE_LEN + 16;
73
74const SHADOW_PARITY_TIERS: [usize; 6] = [4, 8, 16, 32, 64, 128];
76
77const COST_FRACTIONS: [usize; 5] = [20, 10, 5, 2, 1];
82
83const MAX_SHADOW_FRAME_BYTES: usize = 256 * 1024;
85
86fn try_single_fdl(
89 lsbs: &[u8],
90 fdl: usize,
91 parity_len: usize,
92 passphrase: &str,
93) -> Option<Result<PayloadData, StegoError>> {
94 let rs_encoded_len = ecc::rs_encoded_len_with_parity(fdl, parity_len);
95 let rs_bits_needed = rs_encoded_len * 8;
96 if rs_bits_needed > lsbs.len() {
97 return None;
98 }
99
100 let rs_bytes = frame::bits_to_bytes(&lsbs[..rs_bits_needed]);
101 let decoded = match ecc::rs_decode_blocks_with_parity(&rs_bytes, fdl, parity_len) {
102 Ok((data, _stats)) => data,
103 Err(_) => return None,
104 };
105
106 let fr = match parse_shadow_frame(&decoded) {
107 Ok(f) => f,
108 Err(_) => return None,
109 };
110
111 match crypto::decrypt(&fr.ciphertext, passphrase, &fr.salt, &fr.nonce) {
112 Ok(plaintext) => {
113 let len = fr.plaintext_len as usize;
114 if len > plaintext.len() {
115 return None;
116 }
117 Some(payload::decode_payload(&plaintext[..len]))
118 }
119 Err(_) => None,
120 }
121}
122
123fn peek_fdl_from_first_block(
127 lsbs: &[u8],
128 parity_len: usize,
129 max_fdl: usize,
130) -> Option<usize> {
131 let k = 255usize.saturating_sub(parity_len);
132 if k < 2 || lsbs.len() < 255 * 8 {
133 return None;
134 }
135
136 let first_block_bytes = frame::bits_to_bytes(&lsbs[..255 * 8]);
138 let (data, _) = ecc::rs_decode_blocks_with_parity(&first_block_bytes, k, parity_len).ok()?;
139
140 if data.len() < 2 {
141 return None;
142 }
143 let plaintext_len = u16::from_be_bytes([data[0], data[1]]) as usize;
144 let fdl = SHADOW_FRAME_OVERHEAD + plaintext_len;
145
146 if fdl >= k && fdl <= max_fdl {
149 Some(fdl)
150 } else {
151 None
152 }
153}
154
155#[derive(Clone)]
157pub struct ShadowState {
158 pub positions: Vec<CoeffPos>,
160 pub bits: Vec<u8>,
162 pub n_total: usize,
164 pub parity_len: usize,
166 pub frame_data_len: usize,
168 pub frame_bytes: Vec<u8>,
170 pub perm_seed: [u8; 32],
172 pub cost_fraction: usize,
174}
175
176pub fn prepare_shadow(
184 all_y_positions_sorted: &[CoeffPos],
185 shadow_pass: &str,
186 message: &str,
187 files: &[FileEntry],
188 parity_len: usize,
189) -> Result<ShadowState, StegoError> {
190 let payload_bytes = payload::encode_payload(message, files)?;
192
193 let (ciphertext, nonce, salt) = crypto::encrypt(&payload_bytes, shadow_pass)?;
195
196 let frame_bytes = build_shadow_frame(payload_bytes.len(), &salt, &nonce, &ciphertext);
198 let frame_data_len = frame_bytes.len();
199
200 let rs_bytes = ecc::rs_encode_blocks_with_parity(&frame_bytes, parity_len);
202 let rs_bits = frame::bytes_to_bits(&rs_bytes);
203 let n_total = rs_bits.len();
204
205 let perm_seed = crypto::derive_shadow_structural_key(shadow_pass)?;
207 for &fraction in &COST_FRACTIONS {
208 let positions = select_shadow_positions(all_y_positions_sorted, fraction, n_total, &perm_seed);
209 if positions.len() >= n_total {
210 return Ok(ShadowState {
211 positions,
212 bits: rs_bits,
213 n_total,
214 parity_len,
215 frame_data_len,
216 frame_bytes,
217 perm_seed: *perm_seed,
218 cost_fraction: fraction,
219 });
220 }
221 }
222
223 Err(StegoError::MessageTooLarge)
224}
225
226pub fn rebuild_shadow(
234 state: &mut ShadowState,
235 all_y_positions_sorted: &[CoeffPos],
236 new_parity: usize,
237 new_fraction: usize,
238) -> Result<(), StegoError> {
239 let rs_bytes = ecc::rs_encode_blocks_with_parity(&state.frame_bytes, new_parity);
240 let rs_bits = frame::bytes_to_bits(&rs_bytes);
241 let n_total = rs_bits.len();
242
243 let positions = select_shadow_positions(
244 all_y_positions_sorted, new_fraction, n_total, &state.perm_seed,
245 );
246 if positions.len() < n_total {
247 return Err(StegoError::MessageTooLarge);
248 }
249
250 state.positions = positions;
251 state.bits = rs_bits;
252 state.n_total = n_total;
253 state.parity_len = new_parity;
254 state.cost_fraction = new_fraction;
255
256 Ok(())
257}
258
259pub fn embed_shadow_lsb(img: &mut JpegImage, state: &ShadowState) {
264 for (i, pos) in state.positions.iter().enumerate() {
265 if i >= state.n_total {
266 break;
267 }
268 let fi = pos.flat_idx as usize;
269 let coeff = flat_get(img.dct_grid(0), fi);
270 let current_lsb = (coeff.unsigned_abs() & 1) as u8;
271 if current_lsb != state.bits[i] {
272 let modified = nsf5_modify_coefficient(coeff);
273 flat_set(img.dct_grid_mut(0), fi, modified);
274 }
275 }
276}
277
278pub fn verify_shadow(
282 img: &JpegImage,
283 state: &ShadowState,
284 passphrase: &str,
285) -> Result<(), StegoError> {
286 let grid = img.dct_grid(0);
287
288 let lsbs: Vec<u8> = state.positions[..state.n_total].iter().map(|pos| {
290 let coeff = flat_get(grid, pos.flat_idx as usize);
291 (coeff.unsigned_abs() & 1) as u8
292 }).collect();
293 let rs_bytes = frame::bits_to_bytes(&lsbs);
294
295 let (decoded, _stats) = ecc::rs_decode_blocks_with_parity(
297 &rs_bytes, state.frame_data_len, state.parity_len,
298 ).map_err(|_| StegoError::FrameCorrupted)?;
299
300 let fr = parse_shadow_frame(&decoded)?;
302 crypto::decrypt(
303 &fr.ciphertext,
304 passphrase,
305 &fr.salt,
306 &fr.nonce,
307 )?;
308
309 Ok(())
310}
311
312pub fn verify_shadow_decoder_side(
320 img: &JpegImage,
321 state: &ShadowState,
322 passphrase: &str,
323 stego_y_positions_sorted: &[CoeffPos],
324) -> Result<(), StegoError> {
325 let positions = select_shadow_positions(
326 stego_y_positions_sorted, state.cost_fraction, state.n_total, &state.perm_seed,
327 );
328 if positions.len() < state.n_total {
329 return Err(StegoError::FrameCorrupted);
330 }
331
332 let grid = img.dct_grid(0);
333 let lsbs: Vec<u8> = positions[..state.n_total].iter().map(|pos| {
334 let coeff = flat_get(grid, pos.flat_idx as usize);
335 (coeff.unsigned_abs() & 1) as u8
336 }).collect();
337 let rs_bytes = frame::bits_to_bytes(&lsbs);
338
339 let (decoded, _stats) = ecc::rs_decode_blocks_with_parity(
340 &rs_bytes, state.frame_data_len, state.parity_len,
341 ).map_err(|_| StegoError::FrameCorrupted)?;
342
343 let fr = parse_shadow_frame(&decoded)?;
344 crypto::decrypt(
345 &fr.ciphertext,
346 passphrase,
347 &fr.salt,
348 &fr.nonce,
349 )?;
350
351 Ok(())
352}
353
354pub fn shadow_extract(
362 img: &JpegImage,
363 all_y_positions_sorted: &[CoeffPos],
364 passphrase: &str,
365) -> Result<PayloadData, StegoError> {
366 let perm_seed = crypto::derive_shadow_structural_key(passphrase)?;
368
369 let grid = img.dct_grid(0);
370
371 if all_y_positions_sorted.is_empty() {
372 return Err(StegoError::FrameCorrupted);
373 }
374
375 #[cfg(feature = "parallel")]
376 {
377 use rayon::prelude::*;
378
379 let fraction_lsbs: Vec<(usize, Vec<u8>)> = COST_FRACTIONS.par_iter().filter_map(|&fraction| {
381 let pool_size = all_y_positions_sorted.len() / fraction;
382 if pool_size == 0 {
383 return None;
384 }
385 let positions = select_shadow_positions(
386 all_y_positions_sorted, fraction, pool_size, &perm_seed,
387 );
388 if positions.is_empty() {
389 return None;
390 }
391 let lsbs: Vec<u8> = positions.iter().map(|pos| {
392 let coeff = flat_get(grid, pos.flat_idx as usize);
393 (coeff.unsigned_abs() & 1) as u8
394 }).collect();
395 Some((fraction, lsbs))
396 }).collect();
397
398 for (_, lsbs) in &fraction_lsbs {
402 for &parity_len in &SHADOW_PARITY_TIERS {
403 let k = 255usize.saturating_sub(parity_len);
404 if k == 0 { continue; }
405 let max_rs_bytes = lsbs.len() / 8;
406 let max_fdl = compute_max_fdl(max_rs_bytes, parity_len)
407 .min(MAX_SHADOW_FRAME_BYTES);
408 if SHADOW_FRAME_OVERHEAD > max_fdl { continue; }
409
410 if let Some(fdl) = peek_fdl_from_first_block(lsbs, parity_len, max_fdl)
411 && let Some(result) = try_single_fdl(lsbs, fdl, parity_len, passphrase) {
412 return result;
413 }
414 }
415 }
416
417 let mut combos: Vec<(usize, usize, usize)> = Vec::new();
421 for (fi, (_, lsbs)) in fraction_lsbs.iter().enumerate() {
422 for &parity_len in &SHADOW_PARITY_TIERS {
423 let k = 255usize.saturating_sub(parity_len);
424 if k < 2 { continue; }
425 let max_rs_bytes = lsbs.len() / 8;
426 let max_fdl = compute_max_fdl(max_rs_bytes, parity_len)
427 .min(MAX_SHADOW_FRAME_BYTES);
428 let small_max = (k - 1).min(max_fdl);
430 if SHADOW_FRAME_OVERHEAD > small_max { continue; }
431 for fdl in SHADOW_FRAME_OVERHEAD..=small_max {
432 combos.push((fi, parity_len, fdl));
433 }
434 }
435 }
436
437 let result = combos.par_iter().find_map_first(|&(fi, parity_len, fdl)| {
438 let lsbs = &fraction_lsbs[fi].1;
439 try_single_fdl(lsbs, fdl, parity_len, passphrase)
440 });
441
442 match result {
443 Some(ok_or_err) => ok_or_err,
444 None => Err(StegoError::FrameCorrupted),
445 }
446 }
447
448 #[cfg(not(feature = "parallel"))]
449 {
450 let mut last_err = StegoError::FrameCorrupted;
451
452 for &fraction in &COST_FRACTIONS {
455 let pool_size = all_y_positions_sorted.len() / fraction;
456 if pool_size == 0 {
457 continue;
458 }
459
460 let positions = select_shadow_positions(
462 all_y_positions_sorted, fraction, pool_size, &perm_seed,
463 );
464 if positions.is_empty() {
465 continue;
466 }
467
468 let all_lsbs: Vec<u8> = positions.iter().map(|pos| {
470 let coeff = flat_get(grid, pos.flat_idx as usize);
471 (coeff.unsigned_abs() & 1) as u8
472 }).collect();
473
474 if let Some(result) = try_extract_with_lsbs(
476 &all_lsbs, passphrase, &mut last_err,
477 ) {
478 return result;
479 }
480 }
481
482 Err(last_err)
483 }
484}
485
486#[cfg(not(feature = "parallel"))]
490fn try_extract_with_lsbs(
491 all_lsbs: &[u8],
492 passphrase: &str,
493 _last_err: &mut StegoError,
494) -> Option<Result<PayloadData, StegoError>> {
495 for &parity_len in &SHADOW_PARITY_TIERS {
496 let k = 255usize.saturating_sub(parity_len);
497 if k < 2 { continue; }
498
499 let max_rs_bytes = all_lsbs.len() / 8;
500 let max_fdl = compute_max_fdl(max_rs_bytes, parity_len)
501 .min(MAX_SHADOW_FRAME_BYTES);
502 if SHADOW_FRAME_OVERHEAD > max_fdl { continue; }
503
504 if let Some(fdl) = peek_fdl_from_first_block(all_lsbs, parity_len, max_fdl) {
506 if let Some(result) = try_single_fdl(all_lsbs, fdl, parity_len, passphrase) {
507 return Some(result);
508 }
509 }
510
511 let small_max = (k - 1).min(max_fdl);
513 if SHADOW_FRAME_OVERHEAD > small_max { continue; }
514 for fdl in SHADOW_FRAME_OVERHEAD..=small_max {
515 if let Some(result) = try_single_fdl(all_lsbs, fdl, parity_len, passphrase) {
516 return Some(result);
517 }
518 }
519 }
520
521 None
522}
523
524pub fn shadow_capacity(y_nzac: usize) -> usize {
528 if y_nzac == 0 {
529 return 0;
530 }
531
532 let parity_len = SHADOW_PARITY_TIERS[0]; let available_rs_bytes = y_nzac / 8;
534
535 let k = 255 - parity_len;
536 if k == 0 || available_rs_bytes == 0 {
537 return 0;
538 }
539
540 let full_blocks = available_rs_bytes / 255;
541 let remainder_bytes = available_rs_bytes % 255;
542
543 let mut max_frame_bytes = full_blocks * k;
544 if remainder_bytes > parity_len {
545 max_frame_bytes += remainder_bytes - parity_len;
546 }
547
548 max_frame_bytes.saturating_sub(SHADOW_FRAME_OVERHEAD)
549}
550
551fn build_shadow_frame(
557 plaintext_len: usize,
558 salt: &[u8; SALT_LEN],
559 nonce: &[u8; NONCE_LEN],
560 ciphertext: &[u8],
561) -> Vec<u8> {
562 assert!(plaintext_len <= u16::MAX as usize, "shadow frame plaintext exceeds u16::MAX");
563 let mut fr = Vec::with_capacity(SHADOW_FRAME_OVERHEAD + plaintext_len);
564 fr.extend_from_slice(&(plaintext_len as u16).to_be_bytes());
565 fr.extend_from_slice(salt);
566 fr.extend_from_slice(nonce);
567 fr.extend_from_slice(ciphertext);
568 fr
569}
570
571struct ParsedShadowFrame {
573 plaintext_len: u16,
574 salt: [u8; SALT_LEN],
575 nonce: [u8; NONCE_LEN],
576 ciphertext: Vec<u8>,
577}
578
579fn parse_shadow_frame(data: &[u8]) -> Result<ParsedShadowFrame, StegoError> {
581 if data.len() < SHADOW_FRAME_OVERHEAD {
582 return Err(StegoError::FrameCorrupted);
583 }
584
585 let plaintext_len = u16::from_be_bytes([data[0], data[1]]);
586 let expected_len = SHADOW_FRAME_OVERHEAD + plaintext_len as usize;
587 if data.len() < expected_len {
588 return Err(StegoError::FrameCorrupted);
589 }
590
591 let mut salt = [0u8; SALT_LEN];
592 salt.copy_from_slice(&data[2..2 + SALT_LEN]);
593
594 let mut nonce = [0u8; NONCE_LEN];
595 nonce.copy_from_slice(&data[2 + SALT_LEN..2 + SALT_LEN + NONCE_LEN]);
596
597 let ciphertext = data[2 + SALT_LEN + NONCE_LEN..expected_len].to_vec();
598
599 Ok(ParsedShadowFrame {
600 plaintext_len,
601 salt,
602 nonce,
603 ciphertext,
604 })
605}
606
607fn select_shadow_positions(
620 cost_sorted_positions: &[CoeffPos],
621 fraction: usize,
622 n_total: usize,
623 seed: &[u8; 32],
624) -> Vec<CoeffPos> {
625 use rand::{RngCore, SeedableRng};
626 use rand_chacha::ChaCha20Rng;
627
628 let pool_size = cost_sorted_positions.len() / fraction;
630 if pool_size == 0 {
631 return Vec::new();
632 }
633 let pool = &cost_sorted_positions[..pool_size];
634
635 let mut rng = ChaCha20Rng::from_seed(*seed);
637 let mut candidates: Vec<(u64, CoeffPos)> = pool.iter().map(|p| {
638 rng.set_word_pos(p.flat_idx as u128 * 2);
639 let priority = rng.next_u64();
640 (priority, p.clone())
641 }).collect();
642 candidates.sort_by_key(|(priority, _)| *priority);
643
644 candidates.into_iter().map(|(_, p)| p).take(n_total).collect()
645}
646
647fn compute_max_fdl(max_rs_bytes: usize, parity_len: usize) -> usize {
649 let k = 255usize.saturating_sub(parity_len);
650 if k == 0 || max_rs_bytes == 0 {
651 return 0;
652 }
653 let full_blocks = max_rs_bytes / 255;
654 let remainder = max_rs_bytes % 255;
655 let mut max_data = full_blocks * k;
656 if remainder > parity_len {
657 max_data += remainder - parity_len;
658 }
659 max_data
660}
661
662#[cfg(test)]
663mod tests {
664 use super::*;
665
666 #[test]
667 fn shadow_frame_roundtrip() {
668 let salt = [1u8; SALT_LEN];
669 let nonce = [2u8; NONCE_LEN];
670 let ciphertext = vec![0xAA; 20]; let fr = build_shadow_frame(4, &salt, &nonce, &ciphertext);
672 let parsed = parse_shadow_frame(&fr).unwrap();
673 assert_eq!(parsed.plaintext_len, 4);
674 assert_eq!(parsed.salt, salt);
675 assert_eq!(parsed.nonce, nonce);
676 assert_eq!(parsed.ciphertext, ciphertext);
677 }
678
679 #[test]
680 fn shadow_capacity_basic() {
681 let cap = shadow_capacity(100_000);
683 assert!(cap > 10_000, "capacity {cap} should be > 10KB for 100K positions");
684 let large_cap = shadow_capacity(3_000_000);
686 assert!(large_cap > 300_000, "capacity {large_cap} should be > 300KB for 3M positions");
687 }
688
689 #[test]
690 fn shadow_capacity_small() {
691 assert_eq!(shadow_capacity(0), 0);
692 assert_eq!(shadow_capacity(7), 0);
694 }
695
696 #[test]
697 fn shadow_capacity_larger_than_chroma_repetition() {
698 let positions = 1_000_000usize;
700 let old_chroma_cap = (positions / 7 / 8).saturating_sub(50);
702 let cap = shadow_capacity(positions);
703 assert!(
704 cap > old_chroma_cap,
705 "shadow ({cap}) should be larger than old chroma R=7 ({old_chroma_cap})"
706 );
707 }
708}