#!/usr/bin/env python3
"""
HLDS/LG Stage 4 firmware decryptor v4.

Stage 4 consumes Stage 3's firmware-only decrypt candidate list and attempts
firmware decryption. It does not decrypt stubs, does not do checksum
verification, and does not perform deep firmware analysis.

Current implemented decrypt path:
    * old_dld_firmware
        32-bit block bit-transposition + repeating XOR, derived by KPA
        against encrypted old-DLD firmware payload padding.

Additional implemented / guarded path, including Stage 3 v4 probe candidates:
    * raw_kpa_or_newstyle_firmware
        best-effort raw/no-DLD KPA-style decrypt attempt. If the input already
        looks plaintext, it is passed through as ALREADY_PLAINTEXT. If KPA
        fails or the output sanity score is too low, the payload is moved to
        Stage 5 review instead of analysis.

Locked project rules:
    * stubs are never encrypted
    * stubs never enter the decryptor
    * checksum verification belongs after analysis locates the checksum routine
    * raw review decrypt probes are not promoted unless sanity thresholds pass
    * promoted raw plaintext/decrypted probes are not duplicated into Stage 5 review
"""

from __future__ import annotations

import argparse
import csv
import datetime as _dt
import hashlib
import json
import math
import os
import shutil
import struct
import sys
from collections import Counter
from pathlib import Path, PureWindowsPath
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple


PRODUCER = "hlds_stage4_decryptor_v4"

DECRYPT_GDR = [2, 8, 17, 24, 30, 23, 0, 13, 5, 31, 20, 12, 18, 10, 6, 26, 21, 27, 11, 16, 14, 28, 7, 1, 22, 3, 19, 9, 29, 15, 25, 4]
ENCRYPT_GDR = [22, 27, 3, 10, 14, 16, 1, 28, 26, 7, 15, 21, 19, 29, 12, 2, 11, 24, 20, 13, 18, 4, 30, 9, 17, 23, 0, 6, 31, 8, 25]

DECRYPT_GCC = [21, 27, 11, 16, 14, 28, 7, 1, 5, 31, 20, 12, 18, 10, 6, 26, 2, 8, 17, 24, 30, 23, 0, 13, 22, 3, 19, 9, 29, 15, 25, 4]
ENCRYPT_GCC = [22, 11, 3, 26, 30, 16, 1, 12, 10, 7, 31, 21, 5, 19, 13, 28, 2, 27, 8, 20, 29, 18, 4, 14, 25, 17, 23, 0, 6, 15, 24, 9]

FAMILIES = {
    "GCC": {"dec": DECRYPT_GCC, "enc": ENCRYPT_GCC},
    "GDR": {"dec": DECRYPT_GDR, "enc": ENCRYPT_GDR},
}


def utc_iso() -> str:
    return _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def rel_to_path(workspace: Path, rel_or_abs: str) -> Path:
    if not rel_or_abs:
        return workspace

    p_native = Path(rel_or_abs)
    if p_native.is_absolute():
        return p_native

    p_win = PureWindowsPath(rel_or_abs)
    if p_win.is_absolute():
        return Path(str(p_win))

    return workspace.joinpath(*p_win.parts)


def relpath_str(path: Path, workspace: Path) -> str:
    try:
        rel = path.resolve().relative_to(workspace.resolve())
    except Exception:
        try:
            rel = Path(os.path.relpath(path, workspace))
        except Exception:
            rel = path
    return str(rel).replace("/", "\\")


def entropy(data: bytes) -> float:
    if not data:
        return 0.0
    counts = Counter(data)
    n = len(data)
    return -sum((c / n) * math.log2(c / n) for c in counts.values())


def ratio(n: int, d: int) -> float:
    return float(n) / float(d) if d else 0.0


def ascii_string_count(data: bytes, min_len: int = 6, max_scan: int = 2 * 1024 * 1024) -> int:
    scan = data[:max_scan]
    count = 0
    run = 0
    for b in scan:
        if 32 <= b <= 126:
            run += 1
        else:
            if run >= min_len:
                count += 1
            run = 0
    if run >= min_len:
        count += 1
    return count


def basic_profile(data: bytes) -> Dict[str, Any]:
    if not data:
        return {
            "size": 0,
            "sha256": hashlib.sha256(b"").hexdigest(),
            "entropy": 0.0,
            "unique_byte_count": 0,
            "zero_ratio": 0.0,
            "ff_ratio": 0.0,
            "ascii_string_count_min6": 0,
            "first_32_hex": "",
            "last_32_hex": "",
            "known_strings": [],
            "most_common_block_hex": "",
            "most_common_block_count": 0,
        }

    counts = Counter(data)
    blocks = Counter(bytes(data[i:i + 4]) for i in range(0, len(data) - 3, 4))
    common_block, common_count = blocks.most_common(1)[0] if blocks else (b"", 0)

    known = []
    sample = data[:min(len(data), 2 * 1024 * 1024)]
    for s in (b"HL-DT-ST", b"HL-DL-ST", b"DVD-ROM", b"CD-RW", b"RW/DVD", b"GCC-", b"GDR-", b"BOOT", b"FLASH"):
        if s in sample:
            known.append(s.decode("ascii", errors="replace"))

    return {
        "size": len(data),
        "sha256": sha256_bytes(data),
        "entropy": round(entropy(data), 4),
        "unique_byte_count": len(counts),
        "zero_ratio": round(ratio(counts.get(0x00, 0), len(data)), 6),
        "ff_ratio": round(ratio(counts.get(0xFF, 0), len(data)), 6),
        "ascii_string_count_min6": ascii_string_count(data, 6),
        "first_32_hex": data[:32].hex(),
        "last_32_hex": data[-32:].hex() if len(data) >= 32 else data.hex(),
        "known_strings": known,
        "most_common_block_hex": common_block.hex(),
        "most_common_block_count": common_count,
    }


def swap_bits(chunk: int, bits_array: Sequence[int]) -> int:
    result = 0
    for i in range(32):
        bit = (chunk & (1 << bits_array[i])) >> bits_array[i]
        result = (result << 1) | bit
    return result


def kpa_crack(data: bytes) -> Tuple[int, int, bytes]:
    blocks = [bytes(data[i:i + 4]) for i in range(0, len(data) - 3, 4)]
    if not blocks:
        raise ValueError("not enough data for 32-bit KPA")
    most_common_block, count = Counter(blocks).most_common(1)[0]
    ciphertext_val = struct.unpack(">I", most_common_block)[0]
    transposed_key = 0xFFFFFFFF ^ ciphertext_val
    return transposed_key, count, most_common_block


def family_hint_from_candidate(candidate: Dict[str, Any]) -> str:
    parts: List[str] = [
        str(candidate.get("normalized_rel", "")),
        str(candidate.get("stage3_record_id", "")),
        str(candidate.get("stage2_record_id", "")),
    ]

    source = candidate.get("source", {}) or {}
    ident = source.get("identity_hint", {}) or {}
    parts.extend([
        str(ident.get("model", "")),
        str(ident.get("version", "")),
    ])

    for key in ("stage1_staged_rel", "stage1_ext"):
        parts.append(str(source.get(key, "")))

    for s in source.get("stage1_sources", []) or []:
        parts.append(str(s.get("root_file_rel", "")))
        parts.append(str(s.get("source_path_rel", "")))

    text = " ".join(parts).upper()

    # Keep this conservative. RAW_ is source/container vocabulary, not a family.
    text = text.replace("RAW_", "")

    if "GDR" in text or "8050L" in text or "8081N" in text or "8082N" in text or "8083N" in text or "8084N" in text or "816" in text:
        return "GDR"
    if "GCC" in text or "424" in text or "4160" in text:
        return "GCC"
    return "UNKNOWN"


def resolve_true_key_for_family(transposed_key: int, family: str) -> Optional[int]:
    if family not in FAMILIES:
        return None
    true_key = swap_bits(transposed_key, FAMILIES[family]["dec"])
    b = [(true_key >> shift) & 0xFF for shift in (24, 16, 8, 0)]
    if b[0] == b[1] == b[2] == b[3]:
        return true_key
    return None


def decrypt_old_dld_payload(data: bytes, family: str, true_key: int) -> bytes:
    out = bytearray(data)
    dec_matrix = FAMILIES[family]["dec"]
    end = len(out) - (len(out) % 4)
    for i in range(0, end, 4):
        chunk = (out[i] << 24) | (out[i + 1] << 16) | (out[i + 2] << 8) | out[i + 3]
        chunk = swap_bits(chunk, dec_matrix) ^ true_key
        out[i] = (chunk >> 24) & 0xFF
        out[i + 1] = (chunk >> 16) & 0xFF
        out[i + 2] = (chunk >> 8) & 0xFF
        out[i + 3] = chunk & 0xFF
    return bytes(out)


def decrypt_raw_kpa_payload(data: bytes, family: str, true_key: int, max_bytes: int) -> bytes:
    """Best-effort raw/no-DLD KPA-style decrypt.

    This intentionally differs from old_dld_firmware by preserving bytes beyond
    max_bytes unchanged, matching the older fleet decryptor behavior where raw
    payloads were decrypted over the first firmware window only.
    """
    out = bytearray(data)
    dec_matrix = FAMILIES[family]["dec"]
    end = min(len(out), max_bytes)
    end -= (end % 4)
    for i in range(0, end, 4):
        chunk = (out[i] << 24) | (out[i + 1] << 16) | (out[i + 2] << 8) | out[i + 3]
        chunk = swap_bits(chunk, dec_matrix) ^ true_key
        out[i] = (chunk >> 24) & 0xFF
        out[i + 1] = (chunk >> 16) & 0xFF
        out[i + 2] = (chunk >> 8) & 0xFF
        out[i + 3] = chunk & 0xFF
    return bytes(out)


def sanity_score_plaintext(data: bytes) -> Dict[str, Any]:
    prof = basic_profile(data)
    score = 0
    evidence: List[str] = []
    warnings: List[str] = []

    if prof["known_strings"]:
        score += 4
        evidence.append("known_ascii_identity_strings_present")
    if prof["ascii_string_count_min6"] >= 10:
        score += 2
        evidence.append("ascii_string_count_ge_10")
    if prof["zero_ratio"] > 0.01 or prof["ff_ratio"] > 0.01:
        score += 1
        evidence.append("padding_or_erased_bytes_present")
    if prof["entropy"] < 7.3:
        score += 1
        evidence.append("entropy_below_packed_expectation")
    if prof["most_common_block_count"] >= 128:
        score += 1
        evidence.append("repeated_word_padding_present")

    if prof["entropy"] > 7.85 and not prof["known_strings"] and prof["ascii_string_count_min6"] < 5:
        score -= 3
        warnings.append("high_entropy_low_text_after_decrypt")

    if score >= 5:
        status = "strong_plaintext_candidate"
    elif score >= 3:
        status = "moderate_plaintext_candidate"
    elif score >= 1:
        status = "weak_plaintext_candidate"
    else:
        status = "poor_plaintext_candidate"

    return {
        "status": status,
        "score": score,
        "evidence": evidence,
        "warnings": warnings,
        "profile": prof,
        "status_is_sanity_check_not_checksum_verification": True,
    }


def choose_family_and_decrypt(data: bytes, candidate: Dict[str, Any]) -> Tuple[str, int, int, int, bytes, Dict[str, Any]]:
    transposed_key, kpa_count, common_block = kpa_crack(data)
    hint = family_hint_from_candidate(candidate)

    order: List[str]
    if hint in FAMILIES:
        order = [hint] + [f for f in FAMILIES if f != hint]
    else:
        order = list(FAMILIES.keys())

    attempts: List[Dict[str, Any]] = []
    for family in order:
        true_key = resolve_true_key_for_family(transposed_key, family)
        if true_key is None:
            attempts.append({
                "family": family,
                "status": "true_key_not_repeating_byte_pattern",
            })
            continue
        plaintext = decrypt_old_dld_payload(data, family, true_key)
        sanity = sanity_score_plaintext(plaintext)
        attempts.append({
            "family": family,
            "status": "attempted",
            "true_xor_key": f"0x{true_key:08X}",
            "sanity_status": sanity["status"],
            "sanity_score": sanity["score"],
            "sanity_evidence": sanity["evidence"],
            "sanity_warnings": sanity["warnings"],
        })

        # If the family hint is explicit, honor it. GCC and GDR can resolve the
        # same repeating true key while producing different plaintext.
        if family == hint:
            return family, true_key, transposed_key, kpa_count, plaintext, {
                "family_hint": hint,
                "family_selection": "selected_from_identity_hint",
                "kpa_common_block_hex": common_block.hex(),
                "attempts": attempts,
                "sanity": sanity,
            }

    # No explicit hint, or hint failed. Pick the highest scoring attempted family.
    scored = [a for a in attempts if a.get("status") == "attempted"]
    if not scored:
        raise ValueError("could not resolve repeating-byte true key for any family")

    best_meta = max(scored, key=lambda a: a.get("sanity_score", -999))
    best_family = best_meta["family"]
    best_true_key = resolve_true_key_for_family(transposed_key, best_family)
    if best_true_key is None:
        raise ValueError("internal error: selected family no longer resolves true key")
    best_plain = decrypt_old_dld_payload(data, best_family, best_true_key)
    best_sanity = sanity_score_plaintext(best_plain)
    return best_family, best_true_key, transposed_key, kpa_count, best_plain, {
        "family_hint": hint,
        "family_selection": "selected_by_plaintext_sanity_score",
        "kpa_common_block_hex": common_block.hex(),
        "attempts": attempts,
        "sanity": best_sanity,
    }


def clean_basename(rel: str) -> str:
    name = PureWindowsPath(rel).name or Path(rel).name
    if name.lower().endswith(".bin"):
        name = name[:-4]
    # Avoid double plaintext suffix if rerun on an already named file.
    name = name.replace("_plaintext", "")
    return name


class Stage4Decryptor:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.workspace = Path(args.workspace).resolve()
        self.candidates_path = Path(args.candidates) if args.candidates else self.workspace / "stage4_decrypt_candidates.json"
        if not self.candidates_path.is_absolute():
            self.candidates_path = self.workspace / self.candidates_path

        self.stage3_manifest_path = Path(args.stage3_manifest) if args.stage3_manifest else self.workspace / "stage3_routing_manifest.json"
        if not self.stage3_manifest_path.is_absolute():
            self.stage3_manifest_path = self.workspace / self.stage3_manifest_path

        self.output_dir = Path(args.output_dir) if args.output_dir else self.workspace / "decrypted_payloads"
        if not self.output_dir.is_absolute():
            self.output_dir = self.workspace / self.output_dir

        self.logs_dir = self.workspace / "logs"
        self.log_path: Optional[Path] = None

        self.candidates_doc: Dict[str, Any] = {}
        self.stage3_doc: Dict[str, Any] = {}
        self.results: List[Dict[str, Any]] = []
        self.inventory_rows: List[Dict[str, Any]] = []
        self.stage5_analysis_inputs: List[Dict[str, Any]] = []
        self.stage5_review_inputs: List[Dict[str, Any]] = []

        self.counts = {
            "candidates_seen": 0,
            "firmware_candidates_seen": 0,
            "stub_candidates_rejected": 0,
            "old_dld_decrypt_attempts": 0,
            "old_dld_decrypted": 0,
            # v4 raw/no-DLD probe accounting. The legacy raw_kpa_* names are
            # retained as aliases for downstream compatibility, but raw_kpa_attempts
            # now means a real transform attempt, not a plaintext passthrough probe.
            "raw_probe_candidates": 0,
            "raw_plaintext_passthrough": 0,
            "raw_decrypt_attempts": 0,
            "raw_decrypt_success": 0,
            "raw_decrypt_low_confidence": 0,
            "raw_kpa_attempts": 0,
            "raw_kpa_decrypted": 0,
            "raw_kpa_low_confidence": 0,
            "unsupported_methods": 0,
            "already_plaintext": 0,
            "failed": 0,
            "missing_input_files": 0,
            "sha256_mismatches": 0,
            "stage3_review_inputs_suppressed_promoted": 0,
            "stage5_analysis_inputs": 0,
            "stage5_review_inputs": 0,
            "errors": 0,
        }

    def log(self, msg: str) -> None:
        print(msg)
        if self.log_path:
            with self.log_path.open("a", encoding="utf-8") as f:
                f.write(msg + "\n")

    def setup(self) -> None:
        self.logs_dir.mkdir(parents=True, exist_ok=True)
        ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_path = self.logs_dir / f"hlds_stage4_decryptor_{ts}.log"

        if self.args.clean and self.output_dir.exists():
            shutil.rmtree(self.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.log(f"[INFO] {PRODUCER} started {utc_iso()}")
        self.log("[INFO] checksum verification is intentionally not performed in Stage 4")

    def load_inputs(self) -> None:
        self.log(f"[INFO] loading decrypt candidates: {self.candidates_path}")
        with self.candidates_path.open("r", encoding="utf-8") as f:
            self.candidates_doc = json.load(f)
        self.counts["candidates_seen"] = len(self.candidates_doc.get("candidates", []))

        if self.stage3_manifest_path.exists():
            self.log(f"[INFO] loading Stage 3 routing manifest: {self.stage3_manifest_path}")
            with self.stage3_manifest_path.open("r", encoding="utf-8") as f:
                self.stage3_doc = json.load(f)
        else:
            self.log(f"[WARN] Stage 3 routing manifest not found: {self.stage3_manifest_path}")

    def run(self) -> None:
        self.setup()
        self.load_inputs()
        self.process_candidates()
        self.build_stage5_handoff()
        self.write_outputs()
        self.print_summary()

    def process_candidates(self) -> None:
        for cand in self.candidates_doc.get("candidates", []):
            self.process_candidate(cand)

    def process_candidate(self, cand: Dict[str, Any]) -> None:
        role = cand.get("role", "unknown")
        method = cand.get("decrypt_method") or cand.get("method")
        rel = cand.get("normalized_rel", "")

        if role == "stub":
            self.counts["stub_candidates_rejected"] += 1
            res = self.result_base(cand, "REJECTED_STUB", "stubs are never encrypted and never enter Stage 4")
            self.results.append(res)
            return

        if role != "firmware":
            self.counts["failed"] += 1
            res = self.result_base(cand, "FAILED", f"candidate role is not firmware: {role}")
            self.results.append(res)
            return

        self.counts["firmware_candidates_seen"] += 1

        in_path = rel_to_path(self.workspace, rel)
        if not in_path.exists():
            self.counts["missing_input_files"] += 1
            res = self.result_base(cand, "MISSING_INPUT", f"input not found: {rel}")
            self.results.append(res)
            self.log(f"[ERROR] missing input: {rel}")
            return

        try:
            data = in_path.read_bytes()
            in_sha = sha256_bytes(data)
            notes: List[str] = []
            if cand.get("sha256") and cand.get("sha256") != in_sha:
                self.counts["sha256_mismatches"] += 1
                notes.append("input_sha256_mismatch")

            if method == "old_dld_firmware":
                self.counts["old_dld_decrypt_attempts"] += 1
                family, true_key, transposed_key, kpa_count, plaintext, meta = choose_family_and_decrypt(data, cand)
                out_sha = sha256_bytes(plaintext)
                base_name = clean_basename(rel)
                out_path = self.output_dir / f"{base_name}_plaintext_{out_sha[:8]}.bin"
                out_path.write_bytes(plaintext)
                self.counts["old_dld_decrypted"] += 1
                sanity = meta.get("sanity", {})
                res = self.result_base(cand, "DECRYPTED", "decrypted old DLD firmware payload")
                res.update({
                    "input_rel": rel,
                    "input_sha256": in_sha,
                    "output_rel": relpath_str(out_path, self.workspace),
                    "output_sha256": out_sha,
                    "size": len(data),
                    "output_size": len(plaintext),
                    "decrypt": {
                        "method": "old_dld_firmware",
                        "cipher_family": family,
                        "family_hint": meta.get("family_hint"),
                        "family_selection": meta.get("family_selection"),
                        "transposed_key": f"0x{transposed_key:08X}",
                        "true_xor_key": f"0x{true_key:08X}",
                        "kpa_common_block_hex": meta.get("kpa_common_block_hex"),
                        "kpa_common_block_count": kpa_count,
                        "attempts": meta.get("attempts", []),
                    },
                    "sanity": sanity,
                    "checksum_verification": {
                        "status": "not_performed",
                        "reason": "checksum routine/range must be recovered during analysis before verification",
                    },
                    "notes": notes,
                })
                self.results.append(res)
                self.add_inventory(res)
                self.log(
                    f"[INFO] decrypted {rel} -> {res['output_rel']} "
                    f"family={family} key=0x{true_key:08X} sanity={sanity.get('status')}"
                )
                return

            if method in ("raw_kpa_or_newstyle_firmware", "newstyle_or_raw_firmware"):
                self.counts["raw_probe_candidates"] += 1
                input_sanity = sanity_score_plaintext(data)
                if input_sanity.get("score", 0) >= self.args.min_raw_plaintext_passthrough_score:
                    self.counts["already_plaintext"] += 1
                    self.counts["raw_plaintext_passthrough"] += 1
                    res = self.result_base(cand, "ALREADY_PLAINTEXT", "raw/no-DLD payload already meets plaintext sanity threshold")
                    res.update({
                        "input_rel": rel,
                        "input_sha256": in_sha,
                        "output_rel": rel,
                        "output_sha256": in_sha,
                        "size": len(data),
                        "output_size": len(data),
                        "decrypt": {"method": "raw_kpa_or_newstyle_firmware", "action": "passthrough_plaintext"},
                        "sanity": input_sanity,
                        "notes": notes + ["raw_input_passed_plaintext_sanity_without_decrypt"],
                    })
                    self.results.append(res)
                    self.add_inventory(res)
                    self.log(f"[INFO] raw/no-DLD already plaintext {rel} sanity={input_sanity.get('status')}")
                    return

                transposed_key, kpa_count, common_block = kpa_crack(data)
                attempts = []
                if transposed_key == 0xFFFFFFFF:
                    self.counts["already_plaintext"] += 1
                    self.counts["raw_plaintext_passthrough"] += 1
                    res = self.result_base(cand, "ALREADY_PLAINTEXT", "raw KPA common block implies plaintext/erased padding")
                    res.update({
                        "input_rel": rel,
                        "input_sha256": in_sha,
                        "output_rel": rel,
                        "output_sha256": in_sha,
                        "size": len(data),
                        "output_size": len(data),
                        "decrypt": {
                            "method": "raw_kpa_or_newstyle_firmware",
                            "action": "passthrough_kpa_plaintext",
                            "transposed_key": f"0x{transposed_key:08X}",
                            "kpa_common_block_hex": common_block.hex(),
                            "kpa_common_block_count": kpa_count,
                        },
                        "sanity": input_sanity,
                        "notes": notes + ["kpa_transposed_key_all_ff_treated_as_plaintext"],
                    })
                    self.results.append(res)
                    self.add_inventory(res)
                    self.log(f"[INFO] raw/no-DLD KPA indicates plaintext {rel}")
                    return

                hint = family_hint_from_candidate(cand)
                self.counts["raw_decrypt_attempts"] += 1
                self.counts["raw_kpa_attempts"] += 1
                order = [hint] + [f for f in FAMILIES if f != hint] if hint in FAMILIES else list(FAMILIES)
                best = None
                for family in order:
                    true_key = resolve_true_key_for_family(transposed_key, family)
                    if true_key is None:
                        attempts.append({"family": family, "status": "true_key_not_repeating_byte_pattern"})
                        continue
                    plain = decrypt_raw_kpa_payload(data, family, true_key, self.args.raw_kpa_max_bytes)
                    sanity = sanity_score_plaintext(plain)
                    attempt = {
                        "family": family,
                        "status": "attempted",
                        "true_xor_key": f"0x{true_key:08X}",
                        "sanity_status": sanity.get("status"),
                        "sanity_score": sanity.get("score"),
                        "sanity_evidence": sanity.get("evidence", []),
                        "sanity_warnings": sanity.get("warnings", []),
                    }
                    attempts.append(attempt)
                    candidate_best = (sanity.get("score", -999), family, true_key, plain, sanity)
                    if best is None or candidate_best[0] > best[0] or (family == hint and candidate_best[0] >= best[0]):
                        best = candidate_best

                if best is None:
                    self.counts["failed"] += 1
                    res = self.result_base(cand, "RAW_KPA_FAILED", "raw KPA could not resolve a repeating-byte true key")
                    res.update({
                        "input_rel": rel,
                        "input_sha256": in_sha,
                        "size": len(data),
                        "decrypt": {
                            "method": "raw_kpa_or_newstyle_firmware",
                            "transposed_key": f"0x{transposed_key:08X}",
                            "kpa_common_block_hex": common_block.hex(),
                            "kpa_common_block_count": kpa_count,
                            "attempts": attempts,
                        },
                        "sanity": input_sanity,
                        "notes": notes + ["raw_kpa_failed_review_required"],
                    })
                    self.results.append(res)
                    self.add_inventory(res)
                    self.log(f"[WARN] raw/no-DLD KPA failed {rel}")
                    return

                score, family, true_key, plaintext, sanity = best
                out_sha = sha256_bytes(plaintext)
                base_name = clean_basename(rel)
                out_path = self.output_dir / f"{base_name}_rawkpa_plaintext_{out_sha[:8]}.bin"
                if score >= self.args.min_raw_decrypt_sanity_score:
                    out_path.write_bytes(plaintext)
                    self.counts["raw_kpa_decrypted"] += 1
                    self.counts["raw_decrypt_success"] += 1
                    res = self.result_base(cand, "DECRYPTED", "decrypted raw/no-DLD firmware with KPA-style method")
                    res.update({
                        "input_rel": rel,
                        "input_sha256": in_sha,
                        "output_rel": relpath_str(out_path, self.workspace),
                        "output_sha256": out_sha,
                        "size": len(data),
                        "output_size": len(plaintext),
                        "decrypt": {
                            "method": "raw_kpa_or_newstyle_firmware",
                            "cipher_family": family,
                            "family_hint": hint,
                            "family_selection": "selected_by_raw_plaintext_sanity_score",
                            "transposed_key": f"0x{transposed_key:08X}",
                            "true_xor_key": f"0x{true_key:08X}",
                            "kpa_common_block_hex": common_block.hex(),
                            "kpa_common_block_count": kpa_count,
                            "raw_kpa_max_bytes": self.args.raw_kpa_max_bytes,
                            "attempts": attempts,
                        },
                        "sanity": sanity,
                        "notes": notes,
                    })
                    self.results.append(res)
                    self.add_inventory(res)
                    self.log(f"[INFO] raw/no-DLD decrypted {rel} -> {res['output_rel']} family={family} key=0x{true_key:08X} sanity={sanity.get('status')}")
                    return

                self.counts["raw_kpa_low_confidence"] += 1
                self.counts["raw_decrypt_low_confidence"] += 1
                res = self.result_base(cand, "RAW_KPA_LOW_CONFIDENCE", "raw KPA output did not meet sanity threshold; held for review")
                res.update({
                    "input_rel": rel,
                    "input_sha256": in_sha,
                    "size": len(data),
                    "candidate_output_sha256": out_sha,
                    "decrypt": {
                        "method": "raw_kpa_or_newstyle_firmware",
                        "cipher_family": family,
                        "family_hint": hint,
                        "transposed_key": f"0x{transposed_key:08X}",
                        "true_xor_key": f"0x{true_key:08X}",
                        "kpa_common_block_hex": common_block.hex(),
                        "kpa_common_block_count": kpa_count,
                        "raw_kpa_max_bytes": self.args.raw_kpa_max_bytes,
                        "attempts": attempts,
                    },
                    "sanity": sanity,
                    "notes": notes + ["raw_kpa_low_confidence_output_not_written"],
                })
                self.results.append(res)
                self.add_inventory(res)
                self.log(f"[WARN] raw/no-DLD decrypt low confidence {rel} sanity={sanity.get('status')} score={score}")
                return

            self.counts["unsupported_methods"] += 1
            res = self.result_base(cand, "UNSUPPORTED_METHOD", f"unsupported decrypt method: {method}")
            self.results.append(res)
            self.log(f"[WARN] unsupported method for {rel}: {method}")
            return

        except Exception as e:
            self.counts["errors"] += 1
            self.counts["failed"] += 1
            res = self.result_base(cand, "FAILED", f"{type(e).__name__}: {e}")
            self.results.append(res)
            self.log(f"[ERROR] failed {rel}: {type(e).__name__}: {e}")

    def result_base(self, cand: Dict[str, Any], status: str, message: str) -> Dict[str, Any]:
        return {
            "stage3_record_id": cand.get("stage3_record_id"),
            "stage2_record_id": cand.get("stage2_record_id"),
            "role": cand.get("role"),
            "original_role": cand.get("original_role"),
            "target_scope": cand.get("target_scope"),
            "container_type": cand.get("container_type"),
            "candidate_rel": cand.get("normalized_rel"),
            "candidate_sha256": cand.get("sha256"),
            "status": status,
            "message": message,
            "source": cand.get("source", {}),
        }

    def add_inventory(self, res: Dict[str, Any]) -> None:
        sanity = res.get("sanity", {})
        profile = sanity.get("profile", {})
        dec = res.get("decrypt", {})
        self.inventory_rows.append({
            "stage3_record_id": res.get("stage3_record_id", ""),
            "stage2_record_id": res.get("stage2_record_id", ""),
            "status": res.get("status", ""),
            "input_rel": res.get("input_rel", ""),
            "output_rel": res.get("output_rel", ""),
            "size": res.get("size", ""),
            "output_sha256": res.get("output_sha256", ""),
            "decrypt_method": dec.get("method", ""),
            "cipher_family": dec.get("cipher_family", ""),
            "transposed_key": dec.get("transposed_key", ""),
            "true_xor_key": dec.get("true_xor_key", ""),
            "kpa_common_block_count": dec.get("kpa_common_block_count", ""),
            "sanity_status": sanity.get("status", ""),
            "sanity_score": sanity.get("score", ""),
            "entropy": profile.get("entropy", ""),
            "zero_ratio": profile.get("zero_ratio", ""),
            "ff_ratio": profile.get("ff_ratio", ""),
            "ascii_string_count_min6": profile.get("ascii_string_count_min6", ""),
            "known_strings": ",".join(profile.get("known_strings", [])),
        })

    def build_stage5_handoff(self) -> None:
        # Decrypted and validated-pass-through firmware outputs are Stage 5 plaintext firmware inputs.
        # v4 also records promoted Stage 3 records so they are not duplicated into review.
        promoted_stage3_ids = set()
        promoted_payload_keys = set()

        for res in self.results:
            if res.get("status") in ("DECRYPTED", "ALREADY_PLAINTEXT"):
                promoted_stage3_ids.add(res.get("stage3_record_id"))
                promoted_payload_keys.add((res.get("candidate_rel") or res.get("input_rel") or res.get("output_rel"), res.get("candidate_sha256") or res.get("input_sha256") or res.get("output_sha256")))
                self.stage5_analysis_inputs.append({
                    "input_kind": "decrypted_firmware" if res.get("status") == "DECRYPTED" else "already_plaintext_firmware",
                    "role": "firmware",
                    "original_role": res.get("original_role"),
                    "effective_role": "firmware",
                    "target_scope": res.get("target_scope"),
                    "container_type": res.get("container_type"),
                    "plaintext_rel": res.get("output_rel"),
                    "plaintext_sha256": res.get("output_sha256"),
                    "source_stage": "stage4",
                    "stage4_status": res.get("status"),
                    "decrypt": res.get("decrypt", {}),
                    "stage4_result_ref": {
                        "stage3_record_id": res.get("stage3_record_id"),
                        "stage2_record_id": res.get("stage2_record_id"),
                        "candidate_rel": res.get("candidate_rel"),
                        "candidate_sha256": res.get("candidate_sha256"),
                    },
                    "source": res.get("source", {}),
                })
            elif res.get("status") in ("RAW_KPA_FAILED", "RAW_KPA_LOW_CONFIDENCE", "UNSUPPORTED_METHOD", "FAILED", "MISSING_INPUT"):
                self.stage5_review_inputs.append({
                    "input_kind": "stage4_review",
                    "role": res.get("role"),
                    "original_role": res.get("original_role"),
                    "effective_role": res.get("role"),
                    "target_scope": res.get("target_scope"),
                    "container_type": res.get("container_type"),
                    "payload_rel": res.get("candidate_rel") or res.get("input_rel"),
                    "payload_sha256": res.get("candidate_sha256") or res.get("input_sha256"),
                    "source_stage": "stage4",
                    "stage4_status": res.get("status"),
                    "stage4_result": res,
                    "source": res.get("source", {}),
                })

        # Include Stage 3 plaintext routes for stubs and raw plaintext firmware.
        # Do not re-add Stage 3 review records for payloads promoted by Stage 4.
        for rec in self.stage3_doc.get("records", []) if self.stage3_doc else []:
            if rec.get("record_type") != "stage3_route":
                continue
            route = rec.get("route", {}) or {}
            route_name = route.get("route", "")
            needs = route.get("needs_decryption")
            rec_id = rec.get("stage3_record_id") or rec.get("record_id")
            rec_key = (rec.get("normalized_rel"), rec.get("sha256"))

            if route_name in ("plaintext_stub_analysis_candidate", "raw_plaintext_firmware_analysis_candidate") and needs is False:
                self.stage5_analysis_inputs.append({
                    "input_kind": route_name,
                    "role": route.get("effective_role") if route.get("effective_role") in ("firmware", "stub") else rec.get("role"),
                    "original_role": rec.get("role"),
                    "effective_role": route.get("effective_role") if route.get("effective_role") in ("firmware", "stub") else rec.get("role"),
                    "target_scope": rec.get("target_scope") or route.get("target_scope"),
                    "container_type": rec.get("container_type"),
                    "plaintext_rel": rec.get("normalized_rel"),
                    "plaintext_sha256": rec.get("sha256"),
                    "source_stage": "stage3",
                    "stage3_route": route,
                    "source": rec.get("source", {}),
                })
            elif "review" in route_name or rec.get("role") == "unknown":
                if rec_id in promoted_stage3_ids or rec_key in promoted_payload_keys:
                    self.counts["stage3_review_inputs_suppressed_promoted"] += 1
                    continue
                self.stage5_review_inputs.append({
                    "input_kind": "review",
                    "role": route.get("effective_role") if route.get("effective_role") in ("firmware", "stub") else rec.get("role"),
                    "original_role": rec.get("role"),
                    "effective_role": route.get("effective_role") if route.get("effective_role") in ("firmware", "stub") else rec.get("role"),
                    "target_scope": rec.get("target_scope") or route.get("target_scope"),
                    "container_type": rec.get("container_type"),
                    "payload_rel": rec.get("normalized_rel"),
                    "payload_sha256": rec.get("sha256"),
                    "source_stage": "stage3",
                    "stage3_route": route,
                    "source": rec.get("source", {}),
                })

        # Deduplicate exact plaintext rels; keep first.
        seen = set()
        deduped = []
        for item in self.stage5_analysis_inputs:
            key = (item.get("plaintext_rel"), item.get("plaintext_sha256"))
            if key not in seen:
                seen.add(key)
                deduped.append(item)
        self.stage5_analysis_inputs = deduped

        # Deduplicate review rels after suppressing promoted payloads.
        seen_review = set()
        deduped_review = []
        for item in self.stage5_review_inputs:
            key = (item.get("payload_rel"), item.get("payload_sha256"), item.get("stage4_status"), item.get("source_stage"))
            if key not in seen_review:
                seen_review.add(key)
                deduped_review.append(item)
        self.stage5_review_inputs = deduped_review

        self.counts["stage5_analysis_inputs"] = len(self.stage5_analysis_inputs)
        self.counts["stage5_review_inputs"] = len(self.stage5_review_inputs)

    def write_outputs(self) -> None:
        manifest = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "candidates_rel": relpath_str(self.candidates_path, self.workspace),
            "stage3_manifest_rel": relpath_str(self.stage3_manifest_path, self.workspace) if self.stage3_manifest_path.exists() else "",
            "output_dir": relpath_str(self.output_dir, self.workspace),
            "log_rel": relpath_str(self.log_path, self.workspace) if self.log_path else "",
            "counts": self.counts,
            "configuration": {
                "raw_kpa_max_bytes": self.args.raw_kpa_max_bytes,
                "min_raw_decrypt_sanity_score": self.args.min_raw_decrypt_sanity_score,
                "min_raw_plaintext_passthrough_score": self.args.min_raw_plaintext_passthrough_score,
            },
            "stage_boundaries": {
                "stage4_does": [
                    "decrypt_firmware_candidates",
                    "old_dld_firmware_decryption",
                    "raw_no_dld_probe_classification",
                    "raw_kpa_or_newstyle_firmware_attempts",
                    "raw_plaintext_passthrough_when_sanity_strong",
                    "record_crypto_material_and_output_hashes",
                    "plaintext_sanity_checks",
                    "prepare_stage5_plaintext_input_list",
                ],
                "stage4_does_not": [
                    "decrypt_stubs",
                    "checksum_verification",
                    "disassembly",
                    "command_handler_analysis",
                    "profile_generation",
                ],
            },
            "locked_rules": {
                "stubs_are_never_encrypted": True,
                "stubs_never_enter_decryptor": True,
                "raw_newstyle_decrypt_attempted_when_routed": True,
                "low_confidence_raw_decrypt_outputs_go_to_review": True,
                "raw_review_probe_candidates_supported": True,
                "promoted_raw_probe_payloads_not_duplicated_in_review": True,
                "checksum_verification_after_analysis": True,
            },
            "results": self.results,
        }

        manifest_path = self.workspace / "stage4_decryption_manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=4), encoding="utf-8")
        self.log(f"[INFO] wrote Stage 4 manifest: {relpath_str(manifest_path, self.workspace)}")

        inv_path = self.workspace / "stage4_decryption_inventory.csv"
        fields = [
            "stage3_record_id", "stage2_record_id", "status", "input_rel", "output_rel",
            "size", "output_sha256", "decrypt_method", "cipher_family", "transposed_key",
            "true_xor_key", "kpa_common_block_count", "sanity_status", "sanity_score",
            "entropy", "zero_ratio", "ff_ratio", "ascii_string_count_min6", "known_strings",
        ]
        with inv_path.open("w", encoding="utf-8", newline="") as f:
            w = csv.DictWriter(f, fieldnames=fields)
            w.writeheader()
            for row in self.inventory_rows:
                w.writerow(row)
        self.log(f"[INFO] wrote Stage 4 inventory: {relpath_str(inv_path, self.workspace)}")

        stage5_doc = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "source_manifest_rel": "stage4_decryption_manifest.json",
            "analysis_input_count": len(self.stage5_analysis_inputs),
            "review_input_count": len(self.stage5_review_inputs),
            "analysis_inputs": self.stage5_analysis_inputs,
            "review_inputs": self.stage5_review_inputs,
        }
        stage5_path = self.workspace / "stage5_analysis_inputs.json"
        stage5_path.write_text(json.dumps(stage5_doc, indent=4), encoding="utf-8")
        self.log(f"[INFO] wrote Stage 5 analysis inputs: {relpath_str(stage5_path, self.workspace)}")

    def print_summary(self) -> None:
        print()
        print("=" * 72)
        print("HLDS STAGE 4 DECRYPTOR SUMMARY")
        print("=" * 72)
        print(f"Workspace              : {self.workspace}")
        print(f"Candidates             : {self.candidates_path}")
        print(f"Output dir             : {self.output_dir}")
        print(f"Manifest               : {self.workspace / 'stage4_decryption_manifest.json'}")
        print(f"Inventory CSV          : {self.workspace / 'stage4_decryption_inventory.csv'}")
        print(f"Stage5 inputs          : {self.workspace / 'stage5_analysis_inputs.json'}")
        print(f"Log file               : {self.log_path}")
        print("-" * 72)
        print(f"Candidates seen        : {self.counts['candidates_seen']}")
        print(f"Firmware candidates    : {self.counts['firmware_candidates_seen']}")
        print(f"Stub candidates reject : {self.counts['stub_candidates_rejected']}")
        print(f"Old DLD attempts       : {self.counts['old_dld_decrypt_attempts']}")
        print(f"Old DLD decrypted      : {self.counts['old_dld_decrypted']}")
        print(f"Raw probe candidates   : {self.counts['raw_probe_candidates']}")
        print(f"Raw plaintext pass-thru: {self.counts['raw_plaintext_passthrough']}")
        print(f"Raw decrypt attempts   : {self.counts['raw_decrypt_attempts']}")
        print(f"Raw decrypt success    : {self.counts['raw_decrypt_success']}")
        print(f"Raw decrypt low conf.  : {self.counts['raw_decrypt_low_confidence']}")
        print(f"Already plaintext      : {self.counts['already_plaintext']}")
        print(f"Stage3 review suppressed: {self.counts['stage3_review_inputs_suppressed_promoted']}")
        print(f"Unsupported methods    : {self.counts['unsupported_methods']}")
        print(f"Missing input files    : {self.counts['missing_input_files']}")
        print(f"SHA mismatches         : {self.counts['sha256_mismatches']}")
        print(f"Failed                 : {self.counts['failed']}")
        print(f"Stage5 analysis inputs : {self.counts['stage5_analysis_inputs']}")
        print(f"Stage5 review inputs   : {self.counts['stage5_review_inputs']}")
        print(f"Errors                 : {self.counts['errors']}")
        print("=" * 72)


def build_arg_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Stage 4 HLDS/LG firmware decryptor. Decrypts firmware only; never decrypts stubs."
    )
    ap.add_argument(
        "--workspace",
        default=r"D:\vscode-workspace\xbox-dvd-tools-v2",
        help="Workspace root containing Stage 3 routing outputs and normalized/decrypted payload folders.",
    )
    ap.add_argument(
        "--candidates",
        default=None,
        help="Path to stage4_decrypt_candidates.json. Defaults to <workspace>\\stage4_decrypt_candidates.json.",
    )
    ap.add_argument(
        "--stage3-manifest",
        default=None,
        help="Path to stage3_routing_manifest.json. Defaults to <workspace>\\stage3_routing_manifest.json.",
    )
    ap.add_argument(
        "--output-dir",
        default=None,
        help="Directory for plaintext outputs. Defaults to <workspace>\\decrypted_payloads.",
    )
    ap.add_argument(
        "--clean",
        action="store_true",
        help="Delete the output dir before decrypting. Does not delete manifests/logs.",
    )
    ap.add_argument(
        "--raw-kpa-max-bytes",
        type=lambda x: int(x, 0),
        default=0x80000,
        help="Maximum bytes to transform for raw/no-DLD KPA attempts. Default: 0x80000.",
    )
    ap.add_argument(
        "--min-raw-decrypt-sanity-score",
        type=int,
        default=5,
        help="Minimum sanity score required to accept raw/no-DLD KPA output as plaintext. Default: 5.",
    )
    ap.add_argument(
        "--min-raw-plaintext-passthrough-score",
        type=int,
        default=7,
        help="Minimum sanity score required to pass a raw/no-DLD candidate through without decrypting. Default: 7.",
    )
    return ap


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_arg_parser().parse_args(argv)
    s = Stage4Decryptor(args)
    try:
        s.run()
        return 0 if s.counts.get("errors", 0) == 0 and s.counts.get("missing_input_files", 0) == 0 else 2
    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
        return 130
    except Exception as e:
        print(f"fatal: {type(e).__name__}: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
