#!/usr/bin/env python3
"""Full ROM dump over rdebug `peek`, compared with the known v2.2 image.

The ROM is mapped at bit `0x3F200000` for 0x1C0000 bytes (1.75 MiB).  A
sequential `peek` of 1 KiB chunks returns the image; the md5 is compared with
the file the project has been using.

    python3 dwt/rom_dump.py --host 100.103.100.103
"""
from __future__ import annotations

import argparse
import hashlib
import sys

import exploit

ROM_BASE_BITS = 0x3F200000
ROM_LEN = 0x1C0000
KNOWN_MD5 = "b46a3e442ccb86134d2356bd70d22fd2"


def main(argv):
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="100.103.100.103")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--timeout", type=float, default=10.0)
    ap.add_argument("--out", default="dwt/evidence/real-unit-rom.bin")
    ap.add_argument("--known", default=KNOWN_MD5)
    a = ap.parse_args(argv[1:])

    rd = exploit.RDebug(a.host, a.port, a.timeout, False)
    rd.connect()
    data = bytearray()
    nchunks = (ROM_LEN + exploit.MAX_XFER - 1) // exploit.MAX_XFER
    done = 0
    try:
        for off in range(0, ROM_LEN, exploit.MAX_XFER):
            n = min(exploit.MAX_XFER, ROM_LEN - off)
            data += rd.peek(ROM_BASE_BITS + off * 8, n)
            done += 1
            if done % 32 == 0 or done == nchunks:
                print(f"\r[.] {done}/{nchunks} chunks", end="", file=sys.stderr,
                      flush=True)
    finally:
        rd.close()
    print(file=sys.stderr)

    with open(a.out, "wb") as fh:
        fh.write(data)
    md5 = hashlib.md5(data).hexdigest()
    print(f"[+] {len(data)} bytes -> {a.out}")
    print(f"[+] md5 {md5}")
    print(f"[+] known {a.known}")
    print("[+] MATCH" if md5 == a.known else "[-] MISMATCH")
    return 0 if md5 == a.known else 1


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv))
    except (exploit.ProtocolError, OSError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(2)