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

Expects the VM restore tree to have already been patched by the Swift firmware pipeline.
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 pyimg4
    Run make fw_patch / make fw_patch_dev / make fw_patch_jb first to patch boot-chain components.
"""

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

from pyimg4 import IM4M, IM4P, IMG4

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

# ══════════════════════════════════════════════════════════════════
# 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"
DEFAULT_IBEC_BOOT_ARGS = b"serial=3 -v debug=0x2014e %s"

# 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"
PATCHER_BINARY_ENV = "VPHONE_PATCHER_BINARY"


# ══════════════════════════════════════════════════════════════════
# 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)


def project_root():
    return os.path.abspath(os.path.join(_SCRIPT_DIR, ".."))


def patcher_binary_path():
    override = os.environ.get(PATCHER_BINARY_ENV, "").strip()
    if override:
        return os.path.abspath(override)
    return os.path.join(project_root(), ".build", "debug", "vphone-cli")


def run_swift_patch_component(component, src_path, output_path):
    """Patch a single component via the Swift FirmwarePatcher CLI."""
    binary = patcher_binary_path()
    if not os.path.isfile(binary):
        print(f"[-] Swift patcher binary not found: {binary}")
        print("    Run: make patcher_build")
        sys.exit(1)

    run(
        [
            binary,
            "patch-component",
            "--component",
            component,
            "--input",
            src_path,
            "--output",
            output_path,
            "--quiet",
        ]
    )


def load_firmware(path):
    """Load firmware file, auto-detecting IM4P vs raw."""
    with open(path, "rb") as f:
        raw = f.read()

    try:
        im4p = IM4P(raw)
        if im4p.payload.compression:
            im4p.payload.decompress()
        return im4p, bytearray(im4p.payload.data), True, raw
    except Exception:
        return None, bytearray(raw), False, raw


def _save_im4p_with_payp(path, fourcc, patched_data, original_raw):
    """Repackage as LZFSE-compressed IM4P and append PAYP from original."""
    with (
        tempfile.NamedTemporaryFile(suffix=".raw", delete=False) as tmp_raw,
        tempfile.NamedTemporaryFile(suffix=".im4p", delete=False) as tmp_im4p,
    ):
        tmp_raw_path = tmp_raw.name
        tmp_im4p_path = tmp_im4p.name
        tmp_raw.write(bytes(patched_data))

    try:
        subprocess.run(
            [
                "pyimg4",
                "im4p",
                "create",
                "-i",
                tmp_raw_path,
                "-o",
                tmp_im4p_path,
                "-f",
                fourcc,
                "--lzfse",
            ],
            check=True,
            capture_output=True,
        )
        output = bytearray(open(tmp_im4p_path, "rb").read())
    finally:
        os.unlink(tmp_raw_path)
        os.unlink(tmp_im4p_path)

    payp_offset = original_raw.rfind(b"PAYP")
    if payp_offset >= 0:
        payp_data = original_raw[payp_offset - 10 :]
        output.extend(payp_data)
        old_len = int.from_bytes(output[2:5], "big")
        output[2:5] = (old_len + len(payp_data)).to_bytes(3, "big")
        print(f"  [+] preserved PAYP ({len(payp_data)} bytes)")

    with open(path, "wb") as f:
        f.write(output)


def find_restore_dir(base_dir):
    for entry in sorted(os.listdir(base_dir)):
        full = os.path.join(base_dir, entry)
        if os.path.isdir(full) and "Restore" in entry:
            return full
    return None


def find_file(base_dir, patterns, label):
    for pattern in patterns:
        matches = sorted(glob.glob(os.path.join(base_dir, pattern)))
        if matches:
            return matches[0]
    print(f"[-] {label} not found. Searched patterns:")
    for pattern in patterns:
        print(f"    {os.path.join(base_dir, pattern)}")
    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}")
    out_path = os.path.join(temp_dir, f"kernelcache.research.vphone600{RAMDISK_KERNEL_SUFFIX}")
    run_swift_patch_component("kernel-base", pristine, out_path)
    print("  [+] base kernel patches applied for ramdisk variant")
    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 the Swift firmware pipeline
    and overwrites it in-place. No hardcoded offsets needed — the ADRP+ADD
    instructions already point to the string location.
    """
    off = data.find(DEFAULT_IBEC_BOOT_ARGS)
    if off < 0:
        print(f"  [-] boot-args: existing string not found ({DEFAULT_IBEC_BOOT_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_path = find_shsh(vm_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")
    txm_patched_raw = os.path.join(temp_dir, "txm.patched.raw")
    im4p_obj, data, _, original_raw = load_firmware(txm_src)
    with open(txm_raw, "wb") as f:
        f.write(bytes(data))
    print(f"  source: {txm_src}")
    print(f"  format: IM4P, {len(data)} bytes")
    run_swift_patch_component("txm", txm_src, txm_patched_raw)
    with open(txm_patched_raw, "rb") as f:
        patched_txm = f.read()
    txm_im4p = os.path.join(temp_dir, "txm.im4p")
    _save_im4p_with_payp(txm_im4p, TXM_FOURCC, patched_txm, 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()