#!/usr/bin/env python3
"""
Patch Ventoy 1.1.18 vtoyjump32.exe / vtoyjump64.exe to skip the "Ventoy disk"
MBR marker check (CheckVentoyDisk), so the ISO can be mounted when the USB
stick uses a third-party MBR (Xorboot) instead of Ventoy's own MBR.

CheckVentoyDisk() in vtoyjump.c does:
    UINT8 check[8] = { 0x56,0x54,0x00,0x47,0x65,0x00,0x48,0x44 };
    ...
    if (memcmp(SectorBuf + 0x190, check, 8) == 0)
        return TRUE;
    return FALSE;

Patches (always return TRUE, code size unchanged):

  vtoyjump64.exe (PE32+, .text VMA 0x140001000, file off 0x400)
    0x140008D0E  0F 94 C3   sete %bl
              ->  B3 01 90   mov $1,%bl ; nop
    (ebx is zeroed at function entry, then copied to eax at 0x140008D11)

  vtoyjump32.exe (PE32, .text VMA 0x00401000, file off 0x400)
    0x00407BEB  75 87              jne  -> false
              ->  90 90              nop nop
    0x00407BF3  0F 85 7B FF FF FF  jne  -> false
              ->  90 90 90 90 90 90  nop x6

Usage:
    python3 patch_vtoyjump.py <dir-with-vtoyjump32.exe-and-vtoyjump64.exe>

Re-running is safe: already patched files are detected and skipped.
"""

import os
import sys

PATCHES = {
    "vtoyjump64.exe": [
        (0x810E, bytes([0x0F, 0x94, 0xC3]), bytes([0xB3, 0x01, 0x90])),
    ],
    "vtoyjump32.exe": [
        (0x6FEB, bytes([0x75, 0x87]), bytes([0x90, 0x90])),
        (0x6FF3, bytes([0x0F, 0x85, 0x7B, 0xFF, 0xFF, 0xFF]),
                 bytes([0x90, 0x90, 0x90, 0x90, 0x90, 0x90])),
    ],
}


def patch_file(path, patches):
    with open(path, "rb") as f:
        data = bytearray(f.read())

    changed = False
    for off, orig, new in patches:
        cur = bytes(data[off:off + len(orig)])
        if cur == new:
            print("  %s @0x%X already patched" % (os.path.basename(path), off))
            continue
        if cur != orig:
            raise SystemExit(
                "  %s @0x%X unexpected bytes: %s (expected %s)" % (
                    os.path.basename(path), off, cur.hex(" "), orig.hex(" ")))
        data[off:off + len(new)] = new
        print("  %s @0x%X patched: %s -> %s" % (
            os.path.basename(path), off, orig.hex(" "), new.hex(" ")))
        changed = True

    if changed:
        with open(path, "wb") as f:
            f.write(data)
    return changed


def main():
    if len(sys.argv) != 2:
        print(__doc__)
        return 1

    d = sys.argv[1]
    rc = 0
    for name, patches in PATCHES.items():
        path = os.path.join(d, name)
        if not os.path.isfile(path):
            print("MISSING: %s" % path)
            rc = 1
            continue
        print("%s (%d bytes)" % (path, os.path.getsize(path)))
        try:
            patch_file(path, patches)
        except SystemExit as e:
            print(e)
            rc = 1

    print("DONE" if rc == 0 else "FAILED")
    return rc


if __name__ == "__main__":
    sys.exit(main())