#!/usr/bin/env python3
"""
Patch Ventoy 1.1.18 initramfs hook for a non-standard layout.

Problem
-------
VtoyTool/vtoydm.c hardcodes:

    int part = 1;            /* ISO partition is assumed to be partition 1 */
    uint64_t offset = 2048;  /* ... starting at sector 2048                */

so `vtoydm -p/-r` always emits device-mapper tables for `<disk>1`.  When the
ISO partition is, say, partition 3, the guest kernel rejects the table
("sda1 too small for target") and Linux live systems fail to boot.

Fix
---
`vtoy_os_param.vtoy_disk_part_id` (packed struct, byte offset 41) already
contains the real partition number.  The hook library gets two small
functions that rewrite the generated tables:

  * chunk table  (-p): <disk>1 + (offset + 2048 - part_start)  ->  <disk>N + (offset - part_start)
  * raw table    (-r): size of <disk>1                        ->  size of <disk>N, device <disk>N

Partition 1 layouts are left untouched.

The script rewrites the newc cpio archives in place (only the changed entry
changes), so names, order, modes and timestamps are preserved:

    ventoy.cpio
      └── ventoy/hook.cpio.xz   (xz)
            └── hook/ventoy-hook-lib.sh   (patched)

Usage:
    python3 patch_ventoy_hook.py <in.cpio> <out.cpio>
"""

import io
import lzma
import sys

MARK = "VTOY-CUSTOM-LAYOUT"

HELPER = '''# %s: vtoydm assumes the ISO partition is number 1 with offset 2048.
# Rewrite generated dm tables when the ISO lives on another partition.
ventoy_dm_fix() {
    vtfile="$1"
    vtdisk="${2#/dev/}"
    vtmode="$3"

    [ -f /ventoy/ventoy_os_param ] || return 0
    vtpart=$($BUSYBOX_PATH/hexdump -s 41 -n 1 -e '1/1 "%%u"' /ventoy/ventoy_os_param 2>/dev/null)
    [ -z "$vtpart" ] && return 0
    [ "$vtpart" = "0" ] && return 0
    [ "$vtpart" = "1" ] && return 0

    case "$vtdisk" in
        nvme*|mmc*|nbd*) vtsep="p" ;;
        *) vtsep="" ;;
    esac
    vtbase="/sys/class/block/${vtdisk}${vtsep}${vtpart}"
    [ -d "$vtbase" ] || return 0
    vtstart=$($CAT "$vtbase/start" 2>/dev/null)
    vtsize=$($CAT "$vtbase/size" 2>/dev/null)
    [ -z "$vtstart" ] && return 0

    vtdev="/dev/${vtdisk}${vtsep}${vtpart}"
    vtold="/dev/${vtdisk}${vtsep}1"

    if [ "$vtmode" = "1" ]; then
        $AWK -v dev="$vtdev" -v old="$vtold" -v st="$vtstart" '
            $3 == "linear" && $4 == old { $4 = dev; $5 = $5 + 2048 - st }
            { print }
        ' "$vtfile" > "${vtfile}.fix"
    else
        $AWK -v dev="$vtdev" -v old="$vtold" -v sz="$vtsize" '
            $3 == "linear" && $4 == old { $1 = 0; $2 = sz; $4 = dev; $5 = 0 }
            { print }
        ' "$vtfile" > "${vtfile}.fix"
    fi
    $BUSYBOX_PATH/mv "${vtfile}.fix" "$vtfile"
    vtlog "dm table fixed: $vtfile dev=$vtdev start=$vtstart part=$vtpart"
}

''' % MARK

ANCHOR_FUNC = "create_ventoy_device_mapper() {\n"

CALLS_AFTER = (
    "    $VTOY_PATH/tool/vtoydm -r -f $VTOY_PATH/ventoy_image_map -d $1 > $VTOY_PATH/ventoy_raw_table    \n"
)

CALLS_ADD = (
    "    # " + MARK + ": fix dm tables for a non-part1 ISO partition\n"
    "    ventoy_dm_fix $VTOY_PATH/ventoy_dm_table \"$1\" 1\n"
    "    ventoy_dm_fix $VTOY_PATH/ventoy_raw_table \"$1\" 0\n"
)

PERSIST_AFTER = (
    "    $VTOY_PATH/tool/vtoydm -p -f $VTOY_PATH/ventoy_persistent_map -d $1 > $VTOY_PATH/persistent_dm_table        \n"
)

PERSIST_ADD = (
    "    ventoy_dm_fix $VTOY_PATH/persistent_dm_table \"$1\" 1\n"
)


def parse_newc(data):
    """Return list of (fields, name, body) preserving header values."""
    entries = []
    pos = 0
    while pos + 110 <= len(data):
        hdr = data[pos:pos + 110]
        if hdr[:6] != b"070701":
            break
        f = [int(hdr[6 + i * 8:14 + i * 8], 16) for i in range(13)]
        namesize = f[11]
        filesize = f[6]
        name_start = pos + 110
        name = data[name_start:name_start + namesize - 1].decode("utf-8", "surrogateescape")
        p = name_start + namesize
        p = (p + 3) & ~3
        body = data[p:p + filesize]
        p += filesize
        p = (p + 3) & ~3
        entries.append({"fields": f, "name": name, "body": body})
        pos = p
        if name == "TRAILER!!!":
            break
    return entries


def build_newc(entries):
    out = bytearray()
    for e in entries:
        f = list(e["fields"])
        name = e["name"].encode("utf-8", "surrogateescape") + b"\0"
        body = e["body"]
        f[6] = len(body)
        f[11] = len(name)
        hdr = b"070701" + b"".join(b"%08X" % (v & 0xFFFFFFFF) for v in f)
        name_pad = (4 - ((110 + len(name)) % 4)) % 4
        out += hdr
        out += name
        out += b"\0" * name_pad
        out += body
        out += b"\0" * ((-len(body)) % 4)
    return bytes(out)


def patch_hook_sh(text):
    if MARK in text:
        return text, False

    if ANCHOR_FUNC not in text:
        raise SystemExit("anchor create_ventoy_device_mapper not found")
    text = text.replace(ANCHOR_FUNC, HELPER + ANCHOR_FUNC, 1)

    if CALLS_AFTER not in text:
        raise SystemExit("anchor vtoydm -r line not found")
    text = text.replace(CALLS_AFTER, CALLS_AFTER + CALLS_ADD, 1)

    if PERSIST_AFTER in text:
        text = text.replace(PERSIST_AFTER, PERSIST_AFTER + PERSIST_ADD, 1)
    else:
        print("  note: persistent mapper anchor not found (skipped)")

    return text, True


def main():
    src, dst = sys.argv[1], sys.argv[2]
    data = open(src, "rb").read()
    entries = parse_newc(data)
    print("ventoy.cpio entries:", [e["name"] for e in entries if e["name"] != "TRAILER!!!"][:20])

    hook_idx = None
    for i, e in enumerate(entries):
        if e["name"] == "ventoy/hook.cpio.xz":
            hook_idx = i
            break
    if hook_idx is None:
        raise SystemExit("ventoy/hook.cpio.xz not found")

    hook_cpio = lzma.decompress(entries[hook_idx]["body"])
    hook_entries = parse_newc(hook_cpio)
    print("hook.cpio entries:", len(hook_entries))

    target = "hook/ventoy-hook-lib.sh"
    idx = None
    for i, e in enumerate(hook_entries):
        if e["name"] == target:
            idx = i
            break
    if idx is None:
        raise SystemExit("%s not found" % target)

    text = hook_entries[idx]["body"].decode("utf-8", "surrogateescape")
    new_text, changed = patch_hook_sh(text)
    if not changed:
        print("already patched")
        return

    hook_entries[idx]["body"] = new_text.encode("utf-8", "surrogateescape")
    new_hook = build_newc(hook_entries)
    entries[hook_idx]["body"] = lzma.compress(new_hook, format=lzma.FORMAT_XZ, preset=6)

    out = build_newc(entries)
    open(dst, "wb").write(out)
    print("written %s (%d bytes)" % (dst, len(out)))


if __name__ == "__main__":
    main()