#!/usr/bin/env python3
# dump_flash.py -- read the entire EFM8 flash over the c2probe CDC link,
# write a .bin, and compare against the stock firmware image.
#
# SAFETY: this script only ever issues Block Read (0x06). It never sends an
# erase/write/lock command, and it never issues a C2 reset between blocks (a
# reset would un-halt the core and break reads). The c2probe firmware itself
# is read-only by design. Flash cannot be harmed.
#
# Workflow: piinit (halt core) -> loop Block Read over 0x0000..0xFFFF in
# 256-byte blocks -> write out/qunexus_device_dump.bin -> read lock byte at
# 0xFBFF -> diff against stock (out/QuNexus_Firmware_v2.2.1.bin, base 0x2400)
# -> focus report on the RECOVERY.md suspect page 0xE800..0xE9FF.
import os, sys, time, select, termios

DEV = os.environ.get("C2PROBE_DEV", "/dev/cu.usbmodem2101")
HERE = os.path.dirname(os.path.abspath(__file__))
OUTDIR = os.path.join(HERE, "..", "out")
STOCK = os.path.join(OUTDIR, "QuNexus_Firmware_v2.2.1.bin")
STOCK_BASE = 0x2400
DUMP_BIN = os.path.join(OUTDIR, "qunexus_device_dump.bin")
FLASH_SIZE = 0x10000          # 64 KB
BLOCK = 256                   # bytes per Block Read (length code 0 = 256)
SUSPECT_LO, SUSPECT_HI = 0xE800, 0xEA00   # RECOVERY.md page (inclusive..exclusive)

# ---- CDC I/O (termios, no pyserial) ----------------------------------------
def open_port():
    fd = os.open(DEV, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
    a = termios.tcgetattr(fd)
    a[2] = termios.CS8 | termios.CLOCAL | termios.CREAD
    a[3] = 0; a[0] = 0; a[1] = 0
    a[4] = termios.B115200; a[5] = termios.B115200
    termios.tcsetattr(fd, termios.TCSANOW, a)
    return fd

def drain(fd, t=0.2):
    r, _, _ = select.select([fd], [], [], t)
    if r:
        try: return os.read(fd, 4096)
        except BlockingIOError: return b""
    return b""

def cmd(fd, line, timeout=10.0):
    os.write(fd, (line + "\n").encode())
    out = b""
    end = time.time() + timeout
    while time.time() < end:
        r, _, _ = select.select([fd], [], [], 0.1)
        if r:
            try: out += os.read(fd, 4096)
            except BlockingIOError: pass
        if b"\n" in out:
            r2, _, _ = select.select([fd], [], [], 0.06)
            if not r2:
                break
    return out.decode("utf-8", "replace").strip()

def piinit(fd):
    r = cmd(fd, "piinit", 12)
    if r != "ok piinit":
        raise RuntimeError(f"piinit failed: {r!r}")
    return True

def block_read(fd, addr, n, retries=3):
    """Issue one Block Read of n bytes (1..256) at addr. Returns bytes or None."""
    lc = n if n < 256 else 0
    for attempt in range(retries):
        r = cmd(fd, f"read {addr:04x} {lc}", 12)
        # expected: "data <addr> <n> <hex>"
        if r.startswith("data "):
            parts = r.split()
            # parts: ['data', 'addr', 'count', hexstr]
            try:
                hexstr = parts[3]
                data = bytes.fromhex(hexstr)
                if len(data) == n:
                    return data
                # short read: pad/record; retry
            except (IndexError, ValueError):
                pass
        if r.startswith("err"):
            # desync likely: re-init and retry
            piinit(fd)
            continue
        # unexpected: retry
    return None

def main():
    os.makedirs(OUTDIR, exist_ok=True)
    fd = open_port()
    drain(fd, 0.4)

    print(f"[piinit] halting EFM8 core...")
    piinit(fd)
    print(f"[piinit] ok")

    buf = bytearray(FLASH_SIZE)
    bad = []          # list of (addr, 'gap')
    t0 = time.time()
    addr = 0
    while addr < FLASH_SIZE:
        n = min(BLOCK, FLASH_SIZE - addr)
        data = block_read(fd, addr, n)
        if data is None:
            print(f"  [FAIL] 0x{addr:04x}..0x{addr+n-1:04x} (gap)")
            bad.append(addr)
            addr += n
            continue
        buf[addr:addr+n] = data
        if (addr % 0x1000) == 0:
            dt = time.time() - t0
            print(f"  0x{addr:04x}  ({dt:.1f}s)")
        addr += n
    dt = time.time() - t0
    print(f"[dump] done in {dt:.1f}s, {len(bad)} gap(s)")

    # lock byte
    lock = block_read(fd, 0xFBFF, 1)
    lockval = lock[0] if lock else None
    print(f"[lock] byte @0xFBFF = {('0x%02x' % lockval) if lockval is not None else 'read FAILED'}")

    with open(DUMP_BIN, "wb") as f:
        f.write(buf)
    print(f"[write] {DUMP_BIN} ({len(buf)} bytes)")

    # ---- compare to stock --------------------------------------------------
    print("\n[compare] vs stock", os.path.basename(STOCK))
    if not os.path.exists(STOCK):
        print("  stock image not found; skipping diff")
    else:
        stock = open(STOCK, "rb").read()
        stock_end = STOCK_BASE + len(stock)          # exclusive
        # compare over the region both cover
        lo = STOCK_BASE
        hi = min(stock_end, FLASH_SIZE)
        diffs = []
        for a in range(lo, hi):
            if buf[a] != stock[a - STOCK_BASE]:
                diffs.append(a)
        print(f"  stock region 0x{lo:04x}..0x{hi-1:04x} ({hi-lo} bytes)")
        print(f"  differing bytes: {len(diffs)}")
        if diffs:
            # group consecutive
            groups = []
            start = prev = diffs[0]
            for a in diffs[1:]:
                if a == prev + 1:
                    prev = a
                else:
                    groups.append((start, prev)); start = prev = a
            groups.append((start, prev))
            print(f"  {len(groups)} contiguous run(s):")
            for s, e in groups[:40]:
                length = e - s + 1
                dev = buf[s:e+1]
                stk = stock[s-STOCK_BASE:e+1-STOCK_BASE]
                print(f"    0x{s:04x}..0x{e:04x} ({length} B) dev={dev.hex()} stock={stk.hex()}")
            if len(groups) > 40:
                print(f"    ... ({len(groups)-40} more)")
        else:
            print("  >>> device app region is BYTE-IDENTICAL to stock <<<")

    # ---- focused report: suspect page --------------------------------------
    print(f"\n[RECOVERY] suspect page 0x{SUSPECT_LO:04x}..0x{SUSPECT_HI-1:04x}")
    page = bytes(buf[SUSPECT_LO:SUSPECT_HI])
    ff = sum(1 for b in page if b == 0xFF)
    nonff = [i for i, b in enumerate(page) if b != 0xFF]
    print(f"  {len(page)} bytes; 0xFF count = {ff}; non-0xFF count = {len(nonff)}")
    if nonff:
        print(f"  non-0xFF offsets (first 40): {[('0x%x'%(SUSPECT_LO+i)) for i in nonff[:40]]}")
    else:
        print("  >>> entire page reads 0xFF (ERASED) -- confirms RECOVERY.md hypothesis <<<")
    # hex dump first 64 bytes of the page
    print("  first 64 bytes:")
    for off in range(0, min(64, len(page)), 16):
        chunk = page[off:off+16]
        print(f"    {SUSPECT_LO+off:04x}: " + " ".join(f"{b:02x}" for b in chunk))

    os.close(fd)
    return 0 if not bad else 1

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