img4-dump 1.1.0

use std::fs;
use std::io::Read;
use std::path::PathBuf;

use anyhow::{Context, Result};
use clap::{ArgAction, Parser, ValueEnum};
use serde::Serialize;

mod parse;
mod crypto;
mod util;

#[cfg(feature = "lzfse")]
mod decompress_lzfse;
#[cfg(feature = "lzss")]
mod decompress_lzss;

/// IMG4 / IM4P / IM4M dumper & decryptor.
#[derive(Parser, Debug)]
#[command(name = "img4-dump", version)]
struct Cli {
    /// Input: .img4, .im4p, or .im4m
    #[arg(value_name = "INPUT", required = true)]
    input: PathBuf,

    /// Output directory; created if missing
    #[arg(short = 'o', long = "outdir", default_value = "img4_dump")]
    outdir: PathBuf,

    /// Overwrite into existing non-empty outdir
    #[arg(short = 'f', long = "force", action = ArgAction::SetTrue)]
    force: bool,

    /// Verbose metadata on stderr
    #[arg(short = 'v', long = "verbose", action = ArgAction::SetTrue)]
    verbose: bool,

    /// Write JSON metadata summary to stdout
    #[arg(long = "json", action = ArgAction::SetTrue)]
    json: bool,

    /// Attempt to decrypt payload (requires --iv and --key, or plaintext KBAG)
    #[arg(long = "decrypt", action = ArgAction::SetTrue)]
    decrypt: bool,

    /// AES mode to use when decrypting (not encoded in KBAG/IM4P)
    #[arg(long = "aes-mode", value_enum, default_value_t = AesMode::Ctr)]
    aes_mode: AesMode,

    /// Hex IV (32 hex chars), overrides KBAG IV (if present)
    #[arg(long = "iv")]
    iv_hex: Option<String>,

    /// Hex Key (32/48/64 hex chars), overrides KBAG Key (if present)
    #[arg(long = "key")]
    key_hex: Option<String>,

    /// If present, write undecoded (still-encrypted) payload too
    #[arg(long = "keep-ciphertext", action = ArgAction::SetTrue)]
    keep_ciphertext: bool,

    /// Try to decompress known formats (lzfse/lzss) after (optional) decryption
    #[arg(long = "decompress", action = ArgAction::SetTrue)]
    decompress: bool,

    /// Dump IM4M (manifest) to outdir
    #[arg(long = "dump-im4m", action = ArgAction::SetTrue)]
    dump_im4m: bool,

    /// Dump IM4R (restore info) to outdir
    #[arg(long = "dump-im4r", action = ArgAction::SetTrue)]
    dump_im4r: bool,

    /// Extract IM4M properties into JSON (img4_dump/im4m.props.json)
    #[arg(long = "dump-im4m-props", action = ArgAction::SetTrue)]
    dump_im4m_props: bool,

    /// Dump IM4M certificate chain (DER + PEM files)
    #[arg(long = "dump-im4m-certs", action = ArgAction::SetTrue)]
    dump_im4m_certs: bool,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
pub enum AesMode {
    Ctr,
    Cbc,
}

#[derive(Debug, Serialize)]
struct Summary {
    container: parse::ContainerKind,
    im4p: Option<parse::Im4pInfo>,
    im4m: Option<parse::Im4mInfoSummary>,
    im4r_len: Option<usize>,
    notes: Vec<String>,
}

fn main() -> Result<()> {
    env_logger::init();

    let cli = Cli::parse();

    // Debug: show parsed CLI options
    eprintln!("DEBUG: Parsed CLI options: {:?}", cli);

    // Read entire input
    eprintln!("DEBUG: Opening input file: {:?}", cli.input);
    let mut f = fs::File::open(&cli.input).with_context(|| format!("open {:?}", cli.input))?;
    let mut bytes = Vec::new();
    f.read_to_end(&mut bytes)?;
    eprintln!("DEBUG: Read {} bytes from input file", bytes.len());

    // Parse container
    eprintln!("DEBUG: Starting container parsing");
    let parsed = parse::parse_img4_like(&bytes)?;
    eprintln!(
        "DEBUG: Parsed container kind: {:?}, im4p: {}, im4m: {}, im4r: {}",
        parsed.kind,
        parsed.im4p.is_some(),
        parsed.im4m.is_some(),
        parsed.im4r.is_some()
    );

    // Prepare output dir
    eprintln!("DEBUG: Ensuring output directory at {:?}", cli.outdir);
    util::ensure_outdir(&cli.outdir, cli.force)?;
    eprintln!("DEBUG: Output directory ready");

    let mut notes = Vec::new();

    // Dump IM4P
    let mut im4p_info = None;
    if let Some(im4p) = &parsed.im4p {
        eprintln!("DEBUG: Processing IM4P payload ({} bytes)", im4p.data.len());

        let base = cli.outdir.join("im4p.bin");
        fs::write(&base, &im4p.data).context("write im4p.bin")?;
        if cli.verbose {
            eprintln!("wrote {:?}", base);
        }

        // If present, persist KBAG DER blob
        if let Some(kbag_raw) = &im4p.kbag_der {
            eprintln!("DEBUG: KBAG detected, writing DER blob ({} bytes)", kbag_raw.len());
            let p = cli.outdir.join("im4p.kbag.der");
            fs::write(&p, kbag_raw)?;
            if cli.verbose {
                eprintln!("wrote {:?}", p);
            }
        }

        // Optionally keep ciphertext copy (if decryption requested)
        if cli.decrypt && cli.keep_ciphertext {
            eprintln!("DEBUG: Keeping ciphertext copy of IM4P");
            let p = cli.outdir.join("im4p.ciphertext");
            fs::write(&p, &im4p.data)?;
        }

        // Decrypt if requested and IV+Key available (from CLI or KBAG)
        let mut clear = None;
        if cli.decrypt {
            eprintln!("DEBUG: Decryption requested, resolving IV and key");
            let (iv, key) = util::resolve_iv_key(&cli, im4p)?;
            eprintln!("DEBUG: IV resolved ({} bytes), Key resolved ({} bytes)", iv.len(), key.len());

            let mode = cli.aes_mode;
            eprintln!("DEBUG: Using AES mode: {:?}", mode);
            let dec = crypto::decrypt_aes(&im4p.data, &iv, &key, mode)
                .with_context(|| "AES decryption failed (check mode/IV/Key)")?;
            eprintln!("DEBUG: Decryption succeeded, plaintext size {} bytes", dec.len());

            // Validate decryption
            let (valid, detected) = util::validate_decryption(&dec);
            if valid {
                if let Some(fmt) = detected {
                    eprintln!("DEBUG: Decryption validation: OK (detected: {})", fmt);
                    if cli.verbose {
                        eprintln!("Decryption appears valid: {}", fmt);
                    }
                }
            } else {
                let reason = detected.unwrap_or_else(|| "unknown".into());
                eprintln!("WARNING: Decryption validation FAILED: {}", reason);
                eprintln!("WARNING: The output may be garbage (wrong key/IV/mode?)");
                if cli.verbose {
                    eprintln!("TIP: Try different --aes-mode (cbc/ctr) or verify key/IV");
                }
                notes.push(format!("decryption validation failed: {}", reason));
            }

            let out = cli.outdir.join("im4p.decrypted");
            fs::write(&out, &dec)?;
            if cli.verbose {
                eprintln!("wrote {:?}", out);
            }
            clear = Some(dec);
        }

        // Optionally decompress decrypted (preferred) or raw
        if cli.decompress {
            eprintln!("DEBUG: Decompression requested");
            let src: &[u8] = if let Some(ref d) = clear { d } else { &im4p.data };
            eprintln!("DEBUG: Decompression source size {} bytes", src.len());

            match util::try_decompress(src) {
                Ok(Some((name, dec))) => {
                    eprintln!("DEBUG: Decompression succeeded, generated {} ({} bytes)", name, dec.len());
                    let p = cli.outdir.join(name);
                    fs::write(&p, &dec)?;
                    if cli.verbose {
                        eprintln!("wrote {:?}", p);
                    }
                }
                Ok(None) => {
                    eprintln!("DEBUG: No known compression detected");
                    notes.push("no known compression detected".into())
                }
                Err(e) => {
                    eprintln!("DEBUG: Decompression error: {}", e);
                    notes.push(format!("decompress error: {e}"))
                }
            }
        }

        im4p_info = Some(parse::Im4pInfo {
            r#type: im4p.r#type.clone(),
            description: im4p.description.clone(),
            data_len: im4p.data.len(),
            kbag: im4p.kbag_summary.clone(),
        });
        eprintln!("DEBUG: IM4P info recorded");
    }

    // Dump IM4M
    let mut im4m_summary = None;
    if let Some(im4m) = &parsed.im4m {
        eprintln!("DEBUG: IM4M present, size {} bytes", im4m.raw.len());

        if cli.dump_im4m {
            let p = cli.outdir.join("im4m.der");
            fs::write(&p, &im4m.raw)?;
            if cli.verbose {
                eprintln!("wrote {:?}", p);
            }
        }

        // Dump IM4M certificate chain (DER + PEM)
        if cli.dump_im4m_certs {
            let certs = parse::extract_im4m_cert_chain(&im4m.raw)?;
            for (i, der) in certs.iter().enumerate() {
                let der_path = cli.outdir.join(format!("im4m.cert.{i}.der"));
                fs::write(&der_path, der)?;
                let pem_path = cli.outdir.join(format!("im4m.cert.{i}.pem"));
                write_pem_certificate(&pem_path, der)?;
                if cli.verbose {
                    eprintln!("wrote {:?} and {:?}", der_path, pem_path);
                }
            }
        }

        // Dump IM4M properties (typed) if requested
        if cli.dump_im4m_props {
            let props = parse::extract_im4m_properties(&im4m.raw)?;
            let p = cli.outdir.join("im4m.props.json");
            fs::write(&p, serde_json::to_vec_pretty(&props)?)?;
            if cli.verbose {
                eprintln!("wrote {:?}", p);
            }
        }

        im4m_summary = Some(parse::summarize_im4m(im4m)?);
        eprintln!("DEBUG: IM4M summary generated");
    }


    // Dump IM4R
    let mut im4r_len = None;
    if let Some(im4r) = &parsed.im4r {
        eprintln!("DEBUG: IM4R present, length {} bytes", im4r.len());
        if cli.dump_im4r {
            let p = cli.outdir.join("im4r.der");
            fs::write(&p, im4r)?;
            if cli.verbose {
                eprintln!("wrote {:?}", p);
            }
        }
        im4r_len = Some(im4r.len());
    }

    let summary = Summary {
        container: parsed.kind,
        im4p: im4p_info,
        im4m: im4m_summary,
        im4r_len,
        notes,
    };

    eprintln!("DEBUG: Final summary prepared: {:#?}", summary);

    if cli.json {
        println!("{}", serde_json::to_string_pretty(&summary)?);
    } else if cli.verbose {
        eprintln!("{:#?}", summary);
    }

    Ok(())
}

fn write_pem_certificate(path: &std::path::Path, der: &[u8]) -> Result<()> {
    use base64::engine::general_purpose::STANDARD;
    use base64::Engine;

    let b64 = STANDARD.encode(der);
    let mut out = String::with_capacity(b64.len() * 4 / 3 + 128);
    out.push_str("-----BEGIN CERTIFICATE-----\n");
    for chunk in b64.as_bytes().chunks(64) {
        out.push_str(std::str::from_utf8(chunk).unwrap());
        out.push('\n');
    }
    out.push_str("-----END CERTIFICATE-----\n");
    fs::write(path, out)?;
    Ok(())
}