#!/usr/bin/env python3
"""
build_ramdisk.py — Build a signed SSH ramdisk for vphone600.

Expects firmware already patched by patch_firmware.py.
Extracts patched components, signs with SHSH, and builds SSH ramdisk.

Usage:
    python3 build_ramdisk.py [vm_directory]

Directory structure:
    ./shsh/              — SHSH blobs (auto-discovered)
    ./ramdisk_input/     — Tools and SSH resources (auto-setup from CFW)
    ./ramdisk_builder_temp/ — Intermediate .raw files (cleaned up)
    ./Ramdisk/           — Final signed IMG4 output

Prerequisites:
    pip install keystone-engine capstone pyimg4
    Run patch_firmware.py first to patch boot-chain components.
"""

import gzip
import glob
import os
import plistlib
import shutil
import subprocess
import sys

# Ensure sibling modules (patch_firmware) are importable when run from any CWD
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if _SCRIPT_DIR not in sys.path:
    sys.path.insert(0, _SCRIPT_DIR)

from pyimg4 import IM4M, IM4P, IMG4

from fw_patch import (
    load_firmware,
    _save_im4p_with_payp,
    patch_txm,
    find_restore_dir,
    find_file,
)
from patchers.iboot import IBootPatcher
from patchers.kernel import KernelPatcher

# ══════════════════════════════════════════════════════════════════
# Configuration
# ══════════════════════════════════════════════════════════════════

OUTPUT_DIR = "Ramdisk"
TEMP_DIR = "ramdisk_builder_temp"
INPUT_DIR = "ramdisk_input"
RESTORED_EXTERNAL_PATH = "usr/local/bin/restored_external"
RESTORED_EXTERNAL_SERIAL_MARKER = b"SSHRD_Script Sep 22 2022 18:56:50"

# Ramdisk boot-args
RAMDISK_BOOT_ARGS = b"serial=3 rd=md0 debug=0x2014e -v wdt=-1 %s"

# IM4P fourccs for restore mode
TXM_FOURCC = "trxm"
KERNEL_FOURCC = "rkrn"
RAMDISK_KERNEL_SUFFIX = ".ramdisk"
RAMDISK_KERNEL_IMG4 = "krnl.ramdisk.img4"
SUDO_PASSWORD = os.environ.get("VPHONE_SUDO_PASSWORD", None)

# Files to remove from ramdisk to save space
RAMDISK_REMOVE = [
    "usr/bin/img4tool",
    "usr/bin/img4",
    "usr/sbin/dietappleh13camerad",
    "usr/sbin/dietappleh16camerad",
    "usr/local/bin/wget",
    "usr/local/bin/procexp",
]

# Directories to re-sign in ramdisk
SIGN_DIRS = [
    "usr/local/bin/*",
    "usr/local/lib/*",
    "usr/bin/*",
    "bin/*",
    "usr/lib/*",
    "sbin/*",
    "usr/sbin/*",
    "usr/libexec/*",
]

# Compressed archive of ramdisk_input/ (located next to this script)
INPUT_ARCHIVE = "ramdisk_input.tar.zst"


# ══════════════════════════════════════════════════════════════════
# Setup — extract ramdisk_input/ from zstd archive if needed
# ══════════════════════════════════════════════════════════════════


def setup_input(vm_dir):
    """Ensure ramdisk_input/ exists, extracting from .tar.zst if needed."""
    input_dir = os.path.join(vm_dir, INPUT_DIR)

    if os.path.isdir(input_dir):
        return input_dir

    # Look for archive next to this script, then in vm_dir
    for search_dir in (os.path.join(_SCRIPT_DIR, "resources"), _SCRIPT_DIR, vm_dir):
        archive = os.path.join(search_dir, INPUT_ARCHIVE)
        if os.path.isfile(archive):
            print(f"  Extracting {INPUT_ARCHIVE}...")
            subprocess.run(
                ["tar", "--zstd", "-xf", archive, "-C", vm_dir],
                check=True,
            )
            return input_dir

    print(f"[-] Neither {INPUT_DIR}/ nor {INPUT_ARCHIVE} found.")
    print(f"    Place {INPUT_ARCHIVE} next to this script or in the VM directory.")
    sys.exit(1)


# ══════════════════════════════════════════════════════════════════
# SHSH / signing helpers
# ══════════════════════════════════════════════════════════════════


def find_shsh(shsh_dir):
    """Find first SHSH blob in directory."""
    for ext in ("*.shsh", "*.shsh2"):
        matches = sorted(glob.glob(os.path.join(shsh_dir, ext)))
        if matches:
            return matches[0]
    return None


def extract_im4m(shsh_path, im4m_path):
    """Extract IM4M manifest from SHSH blob (handles gzip-compressed)."""
    raw = open(shsh_path, "rb").read()
    if raw[:2] == b"\x1f\x8b":
        raw = gzip.decompress(raw)
    tmp = shsh_path + ".tmp"
    try:
        open(tmp, "wb").write(raw)
        subprocess.run(
            ["pyimg4", "im4m", "extract", "-i", tmp, "-o", im4m_path],
            check=True,
            capture_output=True,
        )
    finally:
        if os.path.exists(tmp):
            os.remove(tmp)


def sign_img4(im4p_path, img4_path, im4m_path, tag=None):
    """Create IMG4 from IM4P + IM4M using pyimg4 Python API."""
    im4p = IM4P(open(im4p_path, "rb").read())
    if tag:
        im4p.fourcc = tag
    im4m = IM4M(open(im4m_path, "rb").read())
    img4 = IMG4(im4p=im4p, im4m=im4m)
    with open(img4_path, "wb") as f:
        f.write(img4.output())


def run(cmd, **kwargs):
    """Run a command, raising on failure."""
    return subprocess.run(cmd, check=True, **kwargs)


def run_sudo(cmd, **kwargs):
    """Run sudo command non-interactively using VPHONE_SUDO_PASSWORD."""
    if SUDO_PASSWORD:
        return run(
            ["sudo", "-S", *cmd],
            input=f"{SUDO_PASSWORD}\n",
            text=True,
            **kwargs,
        )
    return run(["sudo", *cmd], **kwargs)


def ensure_path_within_vm(path, vm_dir, label):
    """Fail if path escapes vm_dir."""
    vm_real = os.path.realpath(vm_dir)
    path_real = os.path.realpath(path)
    if path_real == vm_real or path_real.startswith(vm_real + os.sep):
        return
    print(f"[-] {label} must be inside VM dir")
    print(f"    VM dir: {vm_real}")
    print(f"    Path:   {path_real}")
    sys.exit(1)


def check_prerequisites():
    """Verify required host tools are available."""
    missing = []
    for tool, pkg in [("gtar", "gnu-tar"), ("ldid", "ldid-procursus"), ("trustcache", "trustcache (make setup_tools)")]:
        if not shutil.which(tool):
            missing.append(f"  {tool:12s} — {pkg}")
    if missing:
        print("[-] Missing required tools:")
        for m in missing:
            print(m)
        print("\n    Run: make setup_tools")
        sys.exit(1)


# ══════════════════════════════════════════════════════════════════
# Firmware extraction and IM4P creation
# ══════════════════════════════════════════════════════════════════


def extract_to_raw(src_path, raw_path):
    """Extract IM4P payload to .raw file. Returns (im4p_obj, data, original_raw)."""
    im4p, data, was_im4p, original_raw = load_firmware(src_path)
    with open(raw_path, "wb") as f:
        f.write(bytes(data))
    return im4p, data, original_raw


def create_im4p_uncompressed(raw_data, fourcc, description, output_path):
    """Create uncompressed IM4P from raw data."""
    new_im4p = IM4P(
        fourcc=fourcc,
        description=description,
        payload=bytes(raw_data),
    )
    with open(output_path, "wb") as f:
        f.write(new_im4p.output())


def build_kernel_img4(kernel_src, output_dir, temp_dir, im4m_path, output_name, temp_tag):
    """Build one signed kernel IMG4 from a kernelcache source file."""
    kc_raw = os.path.join(temp_dir, f"{temp_tag}.raw")
    kc_im4p = os.path.join(temp_dir, f"{temp_tag}.im4p")
    _, data, original_raw = extract_to_raw(kernel_src, kc_raw)
    print(f"  source: {kernel_src}")
    print(f"  format: IM4P, {len(data)} bytes")
    _save_im4p_with_payp(kc_im4p, KERNEL_FOURCC, data, original_raw)
    sign_img4(kc_im4p, os.path.join(output_dir, output_name), im4m_path)
    print(f"  [+] {output_name}")


def _find_pristine_cloudos_kernel():
    """Find a pristine CloudOS vphone600 research kernel from project ipsws/."""
    env_path = os.environ.get("RAMDISK_BASE_KERNEL", "").strip()
    if env_path:
        p = os.path.abspath(env_path)
        if os.path.isfile(p):
            return p
        print(f"  [!] RAMDISK_BASE_KERNEL set but not found: {p}")

    project_root = os.path.abspath(os.path.join(_SCRIPT_DIR, ".."))
    patterns = [
        os.path.join(project_root, "ipsws", "PCC-CloudOS*", "kernelcache.research.vphone600"),
        os.path.join(project_root, "ipsws", "*CloudOS*", "kernelcache.research.vphone600"),
    ]
    for pattern in patterns:
        matches = sorted(glob.glob(pattern))
        if matches:
            return matches[0]
    return None


def derive_ramdisk_kernel_source(kc_src, temp_dir):
    """Get source kernel for krnl.ramdisk.img4 entirely within ramdisk_build flow.

    Priority:
      1) Existing legacy snapshot next to restore kernel (`*.ramdisk`)
      2) Derive from pristine CloudOS kernel by applying base KernelPatcher
    """
    legacy_snapshot = f"{kc_src}{RAMDISK_KERNEL_SUFFIX}"
    if os.path.isfile(legacy_snapshot):
        print(f"  found legacy ramdisk kernel snapshot: {legacy_snapshot}")
        return legacy_snapshot

    pristine = _find_pristine_cloudos_kernel()
    if not pristine:
        print("  [!] pristine CloudOS kernel not found; skipping ramdisk-specific kernel image")
        return None

    print(f"  deriving ramdisk kernel from pristine source: {pristine}")
    im4p_obj, data, was_im4p, original_raw = load_firmware(pristine)
    kp = KernelPatcher(data)
    n = kp.apply()
    print(f"  [+] {n} base kernel patches applied for ramdisk variant")

    out_path = os.path.join(temp_dir, f"kernelcache.research.vphone600{RAMDISK_KERNEL_SUFFIX}")
    if was_im4p and im4p_obj is not None:
        _save_im4p_with_payp(out_path, im4p_obj.fourcc, data, original_raw)
    else:
        with open(out_path, "wb") as f:
            f.write(data)
    return out_path


# ══════════════════════════════════════════════════════════════════
# iBEC boot-args patching
# ══════════════════════════════════════════════════════════════════


def patch_ibec_bootargs(data):
    """Replace normal boot-args with ramdisk boot-args in already-patched iBEC.

    Finds the boot-args string written by patch_firmware.py (via IBootPatcher)
    and overwrites it in-place. No hardcoded offsets needed — the ADRP+ADD
    instructions already point to the string location.
    """
    normal_args = IBootPatcher.BOOT_ARGS
    off = data.find(normal_args)
    if off < 0:
        print(f"  [-] boot-args: existing string not found ({normal_args.decode()!r})")
        return False

    args = RAMDISK_BOOT_ARGS + b"\x00"
    data[off : off + len(args)] = args

    # Zero out any leftover from the previous string
    end = off + len(args)
    while end < len(data) and data[end] != 0:
        data[end] = 0
        end += 1

    print(f'  boot-args -> "{RAMDISK_BOOT_ARGS.decode()}" at 0x{off:X}')
    return True


def patch_restored_external_usbmux_label(mountpoint):
    """Patch restored_external USBMux serial label when RAMDISK_UDID is provided."""
    target_udid = os.environ.get("RAMDISK_UDID", "").strip()
    if not target_udid:
        print("  [*] RAMDISK_UDID not set; keeping default restored_external USBMux label")
        return

    try:
        target_bytes = target_udid.encode("ascii")
    except UnicodeEncodeError:
        print(f"[-] RAMDISK_UDID must be ASCII, got: {target_udid!r}")
        sys.exit(1)

    marker_len = len(RESTORED_EXTERNAL_SERIAL_MARKER)
    if len(target_bytes) > marker_len:
        print(f"[-] RAMDISK_UDID too long for restored_external label ({len(target_bytes)} > {marker_len})")
        print(f"    RAMDISK_UDID={target_udid}")
        sys.exit(1)

    restored_external = os.path.join(mountpoint, RESTORED_EXTERNAL_PATH)
    if not os.path.isfile(restored_external):
        print(f"[-] Missing restored_external for USBMux label patch: {restored_external}")
        sys.exit(1)

    with open(restored_external, "rb") as f:
        data = f.read()

    off = data.find(RESTORED_EXTERNAL_SERIAL_MARKER)
    if off < 0:
        print("[-] Could not find default USBMux serial marker in restored_external")
        sys.exit(1)

    if data.find(RESTORED_EXTERNAL_SERIAL_MARKER, off + 1) >= 0:
        print("[!] Multiple USBMux serial markers found in restored_external; patching first occurrence")

    replacement = target_bytes + (b"\x00" * (marker_len - len(target_bytes)))
    patched = data[:off] + replacement + data[off + marker_len :]

    with open(restored_external, "wb") as f:
        f.write(patched)

    print(f"  [+] Patched restored_external USBMux label to: {target_udid}")


# ══════════════════════════════════════════════════════════════════
# Ramdisk DMG building
# ══════════════════════════════════════════════════════════════════


def build_ramdisk(restore_dir, im4m_path, vm_dir, input_dir, output_dir, temp_dir):
    """Build custom SSH ramdisk from restore DMG."""
    # Read RestoreRamDisk path dynamically from BuildManifest.plist
    bm_path = os.path.join(restore_dir, "BuildManifest.plist")
    with open(bm_path, "rb") as f:
        bm = plistlib.load(f)
    ramdisk_rel = bm["BuildIdentities"][0]["Manifest"]["RestoreRamDisk"]["Info"]["Path"]
    ramdisk_src = os.path.join(restore_dir, ramdisk_rel)
    mountpoint = os.path.join(vm_dir, "SSHRD")
    ramdisk_raw = os.path.join(temp_dir, "ramdisk.raw.dmg")
    ramdisk_custom = os.path.join(temp_dir, "ramdisk1.dmg")
    gtar_bin = shutil.which("gtar")
    ldid_bin = shutil.which("ldid")
    tc_bin = shutil.which("trustcache")

    # Extract base ramdisk
    print("  Extracting base ramdisk...")
    run(
        ["pyimg4", "im4p", "extract", "-i", ramdisk_src, "-o", ramdisk_raw],
        capture_output=True,
    )

    ensure_path_within_vm(mountpoint, vm_dir, "Ramdisk mountpoint")
    os.makedirs(mountpoint, exist_ok=True)

    try:
        # Mount, create expanded copy
        print("  Mounting base ramdisk...")
        run_sudo(
            [
                "hdiutil",
                "attach",
                "-mountpoint",
                mountpoint,
                ramdisk_raw,
                "-nobrowse",
                "-owners",
                "off",
            ]
        )

        print("  Creating expanded ramdisk (254 MB)...")
        run_sudo(
            [
                "hdiutil",
                "create",
                "-size",
                "254m",
                "-imagekey",
                "diskimage-class=CRawDiskImage",
                "-format",
                "UDZO",
                "-fs",
                "APFS",
                "-layout",
                "NONE",
                "-srcfolder",
                mountpoint,
                "-copyuid",
                "root",
                ramdisk_custom,
            ]
        )
        run_sudo(["hdiutil", "detach", "-force", mountpoint])

        # Mount expanded, inject SSH
        print("  Mounting expanded ramdisk...")
        run_sudo(
            [
                "hdiutil",
                "attach",
                "-mountpoint",
                mountpoint,
                ramdisk_custom,
                "-nobrowse",
                "-owners",
                "off",
            ]
        )

        print("  Injecting SSH tools...")
        ssh_tar = os.path.join(input_dir, "ssh.tar.gz")
        run_sudo(
            [gtar_bin, "-x", "--no-overwrite-dir", "-f", ssh_tar, "-C", mountpoint]
        )
        patch_restored_external_usbmux_label(mountpoint)

        # Remove unnecessary files
        for rel_path in RAMDISK_REMOVE:
            full = os.path.join(mountpoint, rel_path)
            if os.path.exists(full):
                os.remove(full)

        # Re-sign Mach-O binaries
        print("  Re-signing Mach-O binaries...")
        signcert = os.path.join(input_dir, "signcert.p12")

        for pattern in SIGN_DIRS:
            for path in glob.glob(os.path.join(mountpoint, pattern)):
                if os.path.isfile(path) and not os.path.islink(path):
                    if (
                        "Mach-O"
                        in subprocess.run(
                            ["file", path],
                            capture_output=True,
                            text=True,
                        ).stdout
                    ):
                        subprocess.run(
                            [ldid_bin, "-S", "-M", f"-K{signcert}", path],
                            capture_output=True,
                        )

        # Fix sftp-server entitlements
        sftp_ents = os.path.join(input_dir, "sftp_server_ents.plist")
        sftp_server = os.path.join(mountpoint, "usr/libexec/sftp-server")
        if os.path.exists(sftp_server):
            run([ldid_bin, f"-S{sftp_ents}", "-M", f"-K{signcert}", sftp_server])

        # Build trustcache
        print("  Building trustcache...")
        tc_raw = os.path.join(temp_dir, "sshrd.raw.tc")
        tc_im4p = os.path.join(temp_dir, "trustcache.im4p")

        run([tc_bin, "create", tc_raw, mountpoint])
        run(
            ["pyimg4", "im4p", "create", "-i", tc_raw, "-o", tc_im4p, "-f", "rtsc"],
            capture_output=True,
        )
        sign_img4(
            tc_im4p,
            os.path.join(output_dir, "trustcache.img4"),
            im4m_path,
        )
        print(f"  [+] trustcache.img4")

    finally:
        run_sudo(["hdiutil", "detach", "-force", mountpoint], capture_output=True)

    # Shrink and sign ramdisk
    run_sudo(["hdiutil", "resize", "-sectors", "min", ramdisk_custom])

    print("  Signing ramdisk...")
    rd_im4p = os.path.join(temp_dir, "ramdisk.im4p")
    run(
        ["pyimg4", "im4p", "create", "-i", ramdisk_custom, "-o", rd_im4p, "-f", "rdsk"],
        capture_output=True,
    )
    sign_img4(
        rd_im4p,
        os.path.join(output_dir, "ramdisk.img4"),
        im4m_path,
    )
    print(f"  [+] ramdisk.img4")


# ══════════════════════════════════════════════════════════════════
# Main
# ══════════════════════════════════════════════════════════════════


def main():
    vm_dir = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else os.getcwd())

    if not os.path.isdir(vm_dir):
        print(f"[-] Not a directory: {vm_dir}")
        sys.exit(1)

    # Find SHSH
    shsh_dir = os.path.join(vm_dir, "shsh")
    shsh_path = find_shsh(shsh_dir)
    if not shsh_path:
        print(f"[-] No SHSH blob found in {shsh_dir}/")
        print("    Place your .shsh file in the shsh/ directory.")
        sys.exit(1)

    # Find restore directory
    restore_dir = find_restore_dir(vm_dir)
    if not restore_dir:
        print(f"[-] No *Restore* directory found in {vm_dir}")
        sys.exit(1)

    # Check host tools
    check_prerequisites()

    # Setup input resources (copy from CFW if needed)
    print(f"[*] Setting up {INPUT_DIR}/...")
    input_dir = setup_input(vm_dir)

    # Create temp and output directories
    temp_dir = os.path.join(vm_dir, TEMP_DIR)
    output_dir = os.path.join(vm_dir, OUTPUT_DIR)
    ensure_path_within_vm(temp_dir, vm_dir, "Temp directory")
    ensure_path_within_vm(output_dir, vm_dir, "Output directory")
    for d in (temp_dir, output_dir):
        if os.path.exists(d):
            shutil.rmtree(d)
        os.makedirs(d)

    print(f"[*] VM directory:      {vm_dir}")
    print(f"[*] Restore directory: {restore_dir}")
    print(f"[*] SHSH blob:         {shsh_path}")

    # Extract IM4M from SHSH
    im4m_path = os.path.join(temp_dir, "vphone.im4m")
    print(f"\n[*] Extracting IM4M from SHSH...")
    extract_im4m(shsh_path, im4m_path)

    # ── 1. iBSS (already patched by patch_firmware.py) ───────────
    print(f"\n{'=' * 60}")
    print(f"  1. iBSS (already patched — extract & sign)")
    print(f"{'=' * 60}")
    ibss_src = find_file(
        restore_dir,
        [
            "Firmware/dfu/iBSS.vresearch101.RELEASE.im4p",
        ],
        "iBSS",
    )
    ibss_raw = os.path.join(temp_dir, "iBSS.raw")
    ibss_im4p = os.path.join(temp_dir, "iBSS.im4p")
    im4p_obj, data, _ = extract_to_raw(ibss_src, ibss_raw)
    create_im4p_uncompressed(data, im4p_obj.fourcc, im4p_obj.description, ibss_im4p)
    sign_img4(
        ibss_im4p,
        os.path.join(output_dir, "iBSS.vresearch101.RELEASE.img4"),
        im4m_path,
    )
    print(f"  [+] iBSS.vresearch101.RELEASE.img4")

    # ── 2. iBEC (already patched — just fix boot-args for ramdisk)
    print(f"\n{'=' * 60}")
    print(f"  2. iBEC (patch boot-args for ramdisk)")
    print(f"{'=' * 60}")
    ibec_src = find_file(
        restore_dir,
        [
            "Firmware/dfu/iBEC.vresearch101.RELEASE.im4p",
        ],
        "iBEC",
    )
    ibec_raw = os.path.join(temp_dir, "iBEC.raw")
    ibec_im4p = os.path.join(temp_dir, "iBEC.im4p")
    im4p_obj, data, _ = extract_to_raw(ibec_src, ibec_raw)
    patch_ibec_bootargs(data)
    create_im4p_uncompressed(data, im4p_obj.fourcc, im4p_obj.description, ibec_im4p)
    sign_img4(
        ibec_im4p,
        os.path.join(output_dir, "iBEC.vresearch101.RELEASE.img4"),
        im4m_path,
    )
    print(f"  [+] iBEC.vresearch101.RELEASE.img4")

    # ── 3. SPTM (sign only) ─────────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  3. SPTM (sign only)")
    print(f"{'=' * 60}")
    sptm_src = find_file(
        restore_dir,
        [
            "Firmware/sptm.vresearch1.release.im4p",
        ],
        "SPTM",
    )
    sign_img4(
        sptm_src,
        os.path.join(output_dir, "sptm.vresearch1.release.img4"),
        im4m_path,
        tag="sptm",
    )
    print(f"  [+] sptm.vresearch1.release.img4")

    # ── 4. DeviceTree (sign only) ────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  4. DeviceTree (sign only)")
    print(f"{'=' * 60}")
    dt_src = find_file(
        restore_dir,
        [
            "Firmware/all_flash/DeviceTree.vphone600ap.im4p",
        ],
        "DeviceTree",
    )
    sign_img4(
        dt_src,
        os.path.join(output_dir, "DeviceTree.vphone600ap.img4"),
        im4m_path,
        tag="rdtr",
    )
    print(f"  [+] DeviceTree.vphone600ap.img4")

    # ── 5. SEP (sign only) ───────────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  5. SEP (sign only)")
    print(f"{'=' * 60}")
    sep_src = find_file(
        restore_dir,
        [
            "Firmware/all_flash/sep-firmware.vresearch101.RELEASE.im4p",
        ],
        "SEP",
    )
    sign_img4(
        sep_src,
        os.path.join(output_dir, "sep-firmware.vresearch101.RELEASE.img4"),
        im4m_path,
        tag="rsep",
    )
    print(f"  [+] sep-firmware.vresearch101.RELEASE.img4")

    # ── 6. TXM (release variant — needs patching) ────────────────
    print(f"\n{'=' * 60}")
    print(f"  6. TXM (patch release variant)")
    print(f"{'=' * 60}")
    txm_src = find_file(
        restore_dir,
        [
            "Firmware/txm.iphoneos.release.im4p",
        ],
        "TXM",
    )
    txm_raw = os.path.join(temp_dir, "txm.raw")
    im4p_obj, data, original_raw = extract_to_raw(txm_src, txm_raw)
    patch_txm(data)
    txm_im4p = os.path.join(temp_dir, "txm.im4p")
    _save_im4p_with_payp(txm_im4p, TXM_FOURCC, data, original_raw)
    sign_img4(
        txm_im4p, os.path.join(output_dir, "txm.img4"), im4m_path
    )
    print(f"  [+] txm.img4")

    # ── 7. Kernelcache (already patched — repack with rkrn) ──────
    print(f"\n{'=' * 60}")
    print(f"  7. Kernelcache (already patched — repack as rkrn)")
    print(f"{'=' * 60}")
    kc_src = find_file(
        restore_dir,
        [
            "kernelcache.research.vphone600",
        ],
        "kernelcache",
    )
    kc_ramdisk_src = derive_ramdisk_kernel_source(kc_src, temp_dir)
    if kc_ramdisk_src:
        print(f"  building {RAMDISK_KERNEL_IMG4} from ramdisk kernel source")
        build_kernel_img4(
            kc_ramdisk_src,
            output_dir,
            temp_dir,
            im4m_path,
            RAMDISK_KERNEL_IMG4,
            "kcache_ramdisk",
        )
        print("  building krnl.img4 from restore kernel")

    build_kernel_img4(
        kc_src,
        output_dir,
        temp_dir,
        im4m_path,
        "krnl.img4",
        "kcache",
    )

    # ── 8. Ramdisk + Trustcache ──────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  8. Ramdisk + Trustcache")
    print(f"{'=' * 60}")
    build_ramdisk(restore_dir, im4m_path, vm_dir, input_dir, output_dir, temp_dir)

    # ── Cleanup ──────────────────────────────────────────────────
    print(f"\n[*] Cleaning up {TEMP_DIR}/...")
    shutil.rmtree(temp_dir, ignore_errors=True)
    sshrd_dir = os.path.join(vm_dir, "SSHRD")
    if os.path.exists(sshrd_dir):
        shutil.rmtree(sshrd_dir, ignore_errors=True)

    # ── Summary ──────────────────────────────────────────────────
    print(f"\n{'=' * 60}")
    print(f"  Ramdisk build complete!")
    print(f"  Output: {output_dir}/")
    print(f"{'=' * 60}")
    for f in sorted(os.listdir(output_dir)):
        size = os.path.getsize(os.path.join(output_dir, f))
        print(f"    {f:45s} {size:>10,} bytes")


if __name__ == "__main__":
    main()
