#!/usr/bin/env python3
"""Convert a redump-style Xbox disc image into an XISO that xemu can boot.

A redump dump is the whole DVD: a video partition first, then the game
partition. xemu only understands the game partition ("XISO"), so a redump ISO
makes it report "Please insert an Xbox disc". This finds the game partition and
copies it out.

    redump_to_xiso.py <image.iso>                 # report only
    redump_to_xiso.py <image.iso> --out game.iso  # extract
    redump_to_xiso.py <image.iso> --in-place      # extract, keep original as
                                                 # <name>.redump.iso.bak
Verify an existing image without converting:
    redump_to_xiso.py <image.iso> --verify
"""
import argparse
import os
import struct
import sys

SECTOR = 2048
MAGIC = b"MICROSOFT*XBOX*MEDIA"
# Offsets (in bytes) where an Xbox game partition is known to start.
# 0 = already an XISO; 0x18300000 = standard redump video-partition size.
CANDIDATE_BASES = (0, 0x18300000, 0xFD90000, 0x2080000)


def find_base(f):
    """Return the byte offset of the game partition, or None."""
    for base in CANDIDATE_BASES:
        f.seek(base + 32 * SECTOR)
        if f.read(20) == MAGIC:
            return base
    return None


def read_root(f, base):
    """Return (root_sector, root_size, entries) for the partition at `base`."""
    f.seek(base + 32 * SECTOR)
    header = f.read(SECTOR)
    root_sec, root_size = struct.unpack("<II", header[20:28])
    f.seek(base + root_sec * SECTOR)
    table = f.read(root_size)

    entries = []
    seen = set()

    def walk(off):
        if off in seen or off * 4 + 14 > len(table):
            return
        seen.add(off)
        left, right, sec, size, _attr, nlen = struct.unpack(
            "<HHIIBB", table[off * 4:off * 4 + 14]
        )
        name = table[off * 4 + 14:off * 4 + 14 + nlen].decode("latin1")
        entries.append((name, sec, size))
        if left:
            walk(left)
        if right:
            walk(right)

    walk(0)
    return root_sec, root_size, entries


def verify(path, base=None):
    """Print a structural check of the game partition. Returns True if sane."""
    total = os.path.getsize(path)
    with open(path, "rb") as f:
        if base is None:
            base = find_base(f)
        if base is None:
            print("  no XDVDFS magic found -- not an Xbox disc image (or corrupt)")
            return False
        root_sec, root_size, entries = read_root(f, base)
        print("  game partition at 0x%X (%d MiB in)" % (base, base // (1024 * 1024)))
        print("  root sector %d, %d entries" % (root_sec, len(entries)))
        ok = True
        for name, sec, size in sorted(entries):
            if base + sec * SECTOR + size > total:
                print("  %12d  %-28s OUT OF RANGE (image truncated)" % (size, name))
                ok = False
            else:
                print("  %12d  %s" % (size, name))
        for name, sec, _size in entries:
            if name.lower() == "default.xbe":
                f.seek(base + sec * SECTOR)
                magic = f.read(4)
                print("  default.xbe magic: %r %s"
                      % (magic, "VALID" if magic == b"XBEH" else "BAD"))
                ok = ok and magic == b"XBEH"
                break
        else:
            print("  no default.xbe in root -- xemu will not boot this")
            ok = False
        return ok


def extract(src, dst, base, chunk=1024 * 1024):
    total = os.path.getsize(src) - base
    done = 0
    with open(src, "rb") as fi, open(dst, "wb") as fo:
        fi.seek(base)
        while True:
            buf = fi.read(chunk)
            if not buf:
                break
            fo.write(buf)
            done += len(buf)
            pct = done * 100 // total if total else 100
            print("\r  %d%% (%d/%d MiB)"
                  % (pct, done // (1024 * 1024), total // (1024 * 1024)),
                  end="", file=sys.stderr, flush=True)
        fo.flush()
        os.fsync(fo.fileno())
    print(file=sys.stderr)
    return done


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("image")
    ap.add_argument("--out", help="write the XISO here")
    ap.add_argument("--in-place", action="store_true",
                    help="replace the image, keeping the original as .redump.iso.bak")
    ap.add_argument("--verify", action="store_true",
                    help="only check the image's structure")
    args = ap.parse_args()

    src = os.path.abspath(args.image)
    print("%s (%d bytes)" % (src, os.path.getsize(src)))

    with open(src, "rb") as f:
        base = find_base(f)

    if base is None:
        print("  no XDVDFS magic at any known offset -- not a usable Xbox image")
        return 1

    if args.verify or (not args.out and not args.in_place):
        print("  format: %s" % ("XISO (xemu-ready)" if base == 0 else "redump (needs conversion)"))
        verify(src, base)
        if base != 0 and not args.verify:
            print("\n  rerun with --in-place (or --out FILE) to convert")
        return 0

    if base == 0:
        print("  already an XISO -- nothing to convert")
        return 0

    dst = args.out or src + ".xiso.part"
    if os.path.exists(dst) and not args.out:
        print("  %s already exists; refusing to overwrite" % dst)
        return 1

    free = os.statvfs(os.path.dirname(dst)).f_bavail * os.statvfs(os.path.dirname(dst)).f_frsize
    need = os.path.getsize(src) - base
    if free < need:
        print("  need %d MiB free, have %d MiB" % (need // 1048576, free // 1048576))
        return 1

    print("  extracting game partition from 0x%X ..." % base)
    extract(src, dst, base)

    print("  verifying %s" % dst)
    if not verify(dst, 0):
        print("  VERIFY FAILED -- leaving %s in place, original untouched" % dst)
        return 1

    if args.in_place:
        backup = os.path.splitext(src)[0] + ".redump.iso.bak"
        os.rename(src, backup)
        os.rename(dst, src)
        print("  done: %s is now an XISO; original kept as %s" % (src, backup))
    else:
        print("  done: %s" % dst)
    return 0


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