img4-dump 3.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;
mod fourcc;
mod formatter;

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

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
pub enum KbagClass {
    Prod,   // kclass=1
    Dev,    // kclass=2
    Any,    // first entry
}

/// 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). Most Apple
    /// images (iBoot/iBEC/iBSS/LLB/SEP/ramdisk/logo) are CBC, hence the default.
    #[arg(long = "aes-mode", value_enum, default_value_t = AesMode::Cbc)]
    aes_mode: AesMode,

    /// One-shot decrypt: with --iv/--key (or a plaintext KBAG), decrypt the IM4P
    /// payload and write "<input>.decrypted" next to the input. Tries CBC then
    /// CTR and keeps whichever validates. Ignores --outdir and the other dumps.
    #[arg(long = "auto", action = ArgAction::SetTrue)]
    auto: bool,

    /// 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,

    /// KBAG selection preference (prod=class 1, dev=class 2, any=first)
    #[arg(long = "kbag-class", value_enum, default_value_t = KbagClass::Prod)]
    kbag_class: KbagClass,

    /// Optional KBAG entry index override (0-based)
    #[arg(long = "kbag-index")]
    kbag_index: Option<usize>,
}

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

#[derive(Debug, Serialize)]
struct OutputFile {
    label: String,
    path: String,
}

#[derive(Debug, Serialize)]
struct Summary {
    container: parse::ContainerKind,
    im4p: Option<parse::Im4pInfo>,
    im4m: Option<parse::Im4mInfoSummary>,
    im4r_len: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    im4r_properties: Option<Vec<parse::TypedIm4mProperty>>,
    notes: Vec<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    output_files: Vec<OutputFile>,
}

fn main() {
    let cli = Cli::parse();

    // Initialize logger based on verbose flag
    if std::env::var("RUST_LOG").is_err() {
        let level = if cli.verbose { "debug" } else { "info" };
        std::env::set_var("RUST_LOG", level);
    }
    env_logger::init();

    let json = cli.json;
    if let Err(e) = run(cli) {
        // In --json mode, emit a machine-readable error object on stdout so a
        // consumer always receives valid JSON; otherwise print the error chain.
        if json {
            let obj = serde_json::json!({ "error": format!("{e:#}") });
            println!("{}", serde_json::to_string_pretty(&obj).unwrap_or_else(|_| "{\"error\":\"unknown\"}".into()));
        } else {
            eprintln!("Error: {e:#}");
        }
        std::process::exit(1);
    }
}

fn run(cli: Cli) -> Result<()> {
    log::debug!("Parsed CLI options: {:?}", cli);

    // Read entire input
    log::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)?;
    log::debug!("Read {} bytes from input file", bytes.len());

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

    // One-shot auto-decrypt mode: focused decrypt -> "<input>.decrypted".
    if cli.auto {
        return run_auto(&cli, &parsed);
    }

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

    let mut notes = Vec::new();
    let mut output_paths = formatter::OutputPaths::default();

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

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

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

        // If present, persist PAYP payload-scoped properties (with-properties variant)
        if let Some(props) = &im4p.payload_properties {
            let p = cli.outdir.join("im4p.payp.json");
            fs::write(&p, serde_json::to_vec_pretty(props)?)?;
            output_paths.add("Payload Properties", p.display().to_string());
            if cli.verbose {
                eprintln!("wrote {:?} ({} properties)", p, props.len());
            }
        }

        // Optionally keep ciphertext copy (if decryption requested)
        if cli.decrypt && cli.keep_ciphertext {
            log::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 {
            log::debug!("Decryption requested, resolving IV and key");
            let (iv, key) = util::resolve_iv_key(&cli, im4p)?;
            log::debug!("IV resolved ({} bytes), Key resolved ({} bytes)", iv.len(), key.len());

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

            // Validate decryption
            let (valid, detected) = util::validate_decryption(&dec);
            if valid {
                if let Some(fmt) = detected {
                    log::debug!("Decryption validation: OK (detected: {})", fmt);
                    if cli.verbose {
                        eprintln!("Decryption appears valid: {}", fmt);
                    }
                }
            } else {
                let reason = detected.unwrap_or_else(|| "unknown".into());
                log::warn!("Decryption validation FAILED: {}", reason);
                log::warn!("The output may be garbage (wrong key/IV/mode?)");
                // Suggest the *other* mode prominently (not just under --verbose):
                // the single most common cause is mode mismatch. iBoot/iBEC/iBSS/
                // LLB/SEP/ramdisk/logo images are AES-CBC, not the CTR default.
                let other = match mode {
                    AesMode::Ctr => "cbc",
                    AesMode::Cbc => "ctr",
                };
                log::warn!(
                    "HINT: this looks like the wrong AES mode — retry with --aes-mode {other} \
                     (most Apple images, e.g. iBoot/iBEC/iBSS/LLB/SEP, are CBC)"
                );
                notes.push(format!("decryption validation failed: {}", reason));
                notes.push(format!("hint: retry with --aes-mode {other}"));
            }

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

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

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

im4p_info = Some(parse::Im4pInfo {
            r#type: im4p.r#type.clone(),
            version: im4p.version.clone(),
            data_len: im4p.data.len(),
            // Redact raw IV/key bytes from the summary: expose only class + lengths.
            kbag: im4p
                .kbag_summary
                .as_ref()
                .map(|v| v.iter().map(parse::KbagEntryInfo::from).collect()),
            compression: im4p.compression.as_ref().map(parse::CompressionInfo::from),
            payload_properties: im4p.payload_properties.clone(),
        });
        log::debug!("IM4P info recorded");
    }

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

        if cli.dump_im4m {
            let p = cli.outdir.join("im4m.der");
            fs::write(&p, &im4m.raw)?;
            output_paths.add("Manifest", p.display().to_string());
            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)?;
                output_paths.add(format!("Certificate {} (DER)", i), der_path.display().to_string());
                
                let pem_path = cli.outdir.join(format!("im4m.cert.{i}.pem"));
                write_pem_certificate(&pem_path, der)?;
                output_paths.add(format!("Certificate {} (PEM)", i), pem_path.display().to_string());
                
                if cli.verbose {
                    eprintln!("wrote {:?} and {:?}", der_path, pem_path);
                }
            }
        }

// Dump IM4M properties (structured: MANP properties + per-image groups)
if cli.dump_im4m_props {
    let manifest = parse::extract_im4m_manifest(&im4m.raw)?;
    let p = cli.outdir.join("im4m.props.json");
    fs::write(&p, serde_json::to_vec_pretty(&manifest)?)?;
    output_paths.add("Manifest Properties", p.display().to_string());
    if cli.verbose {
        eprintln!(
            "wrote {:?} ({} manifest properties, {} image objects)",
            p,
            manifest.manifest_properties.len(),
            manifest.images.len()
        );
    }
}

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


// Dump IM4R
let mut im4r_len = None;
let mut im4r_properties = None;
if let Some(im4r) = &parsed.im4r {
    log::debug!("IM4R present, length {} bytes", im4r.len());

    // Extract all IM4R properties using structured parser
    match parse::extract_im4r_properties(im4r) {
        Ok(props) => {
            // Extract BNCN nonce if present
            if let Some(bncn_prop) = props.iter().find(|p| p.key == "BNCN") {
                if let parse::Im4mPropertyValue::OctetString { value: hex_nonce } = &bncn_prop.value {
                    if let Ok(nonce) = hex::decode(hex_nonce) {
                        let p = cli.outdir.join("im4r.bncn.bin");
                        fs::write(&p, &nonce)?;
                        output_paths.add("IM4R Nonce (BNCN)", p.display().to_string());
                        if cli.verbose {
                            eprintln!("extracted BNCN nonce ({} bytes) -> {:?}", nonce.len(), p);
                        }
                    }
                }
            } else {
                log::debug!("IM4R contains no BNCN property");
            }
            // Dump all IM4R properties to JSON
            let p = cli.outdir.join("im4r.props.json");
            fs::write(&p, serde_json::to_vec_pretty(&props)?)?;
            output_paths.add("IM4R Properties", p.display().to_string());
            if cli.verbose {
                eprintln!("wrote IM4R properties to {:?}", p);
            }
            im4r_properties = Some(props);
        }
        Err(e) => {
            log::warn!("IM4R property parse error: {}", e);
        }
    }

    if cli.dump_im4r {
        let p = cli.outdir.join("im4r.der");
        fs::write(&p, im4r)?;
        output_paths.add("IM4R", p.display().to_string());
        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,
        im4r_properties,
        notes,
        output_files: output_paths
            .files
            .iter()
            .map(|(label, path)| OutputFile { label: label.clone(), path: path.clone() })
            .collect(),
    };

    log::debug!("Final summary prepared: {:#?}", summary);

    if cli.json {
        println!("{}", serde_json::to_string_pretty(&summary)?);
    } else {
        // Format and print clean output
        let formatted = formatter::format_summary(
            summary.container,
            summary.im4p.as_ref(),
            summary.im4m.as_ref(),
            summary.im4r_len,
            &output_paths,
        );
        print!("{}", formatted);
    }

    Ok(())
}

fn mode_str(m: AesMode) -> &'static str {
    match m {
        AesMode::Cbc => "CBC",
        AesMode::Ctr => "CTR",
    }
}

/// One-shot decrypt: resolve IV/key, try the preferred AES mode then the other,
/// keep whichever produces a valid-looking result, and write it next to the
/// input as "<input>.decrypted".
fn run_auto(cli: &Cli, parsed: &parse::Parsed) -> Result<()> {
    let im4p = parsed
        .im4p
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("--auto: input has no IM4P payload to decrypt"))?;

    let (iv, key) = util::resolve_iv_key(cli, im4p)?;

    let other = match cli.aes_mode {
        AesMode::Cbc => AesMode::Ctr,
        AesMode::Ctr => AesMode::Cbc,
    };

    // Try preferred mode first, then the other; prefer one that validates, but
    // keep the first that merely decrypts as a fallback.
    let mut chosen: Option<(AesMode, Vec<u8>)> = None;
    let mut fallback: Option<(AesMode, Vec<u8>)> = None;
    for m in [cli.aes_mode, other] {
        match crypto::decrypt_aes(&im4p.data, &iv, &key, m) {
            Ok(dec) => {
                let (valid, why) = util::validate_decryption(&dec);
                log::debug!("--auto: {} -> valid={} ({:?})", mode_str(m), valid, why);
                if valid {
                    chosen = Some((m, dec));
                    break;
                }
                fallback.get_or_insert((m, dec));
            }
            Err(e) => log::debug!("--auto: {} mode failed: {}", mode_str(m), e),
        }
    }

    let (mode, dec, validated) = match chosen {
        Some((m, d)) => (m, d, true),
        None => {
            let (m, d) = fallback.ok_or_else(|| {
                anyhow::anyhow!(
                    "--auto: decryption failed in every AES mode — check the IV/key, \
                     or the payload may not be encrypted"
                )
            })?;
            (m, d, false)
        }
    };

    // Write "<input>.decrypted" alongside the input file.
    let mut out = cli.input.clone().into_os_string();
    out.push(".decrypted");
    let out = PathBuf::from(out);
    fs::write(&out, &dec).with_context(|| format!("write {:?}", out))?;

    if cli.json {
        let obj = serde_json::json!({
            "mode": mode_str(mode).to_lowercase(),
            "validated": validated,
            "bytes": dec.len(),
            "output": out.display().to_string(),
        });
        println!("{}", serde_json::to_string_pretty(&obj)?);
    } else if validated {
        eprintln!(
            "Decrypted with AES-{} -> {} ({} bytes)",
            mode_str(mode),
            out.display(),
            dec.len()
        );
    } else {
        eprintln!(
            "WARNING: no AES mode produced a valid-looking result; wrote AES-{} output anyway \
             -> {} ({} bytes). Verify the IV/key.",
            mode_str(mode),
            out.display(),
            dec.len()
        );
    }
    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(())
}