#!/usr/bin/env python3
"""fnwatch.py — hardware data-watchpoint (DR0 RW) on the live FIR kernel word.
Catches whoever READS the built kernel during offline render -> consumer RIP."""
import ctypes
import glob
import os
import signal
import struct
import subprocess
import sys
import time

DUR = float(os.environ.get('WATCH_DUR', '20'))
BIN = '/home/m/re-tools/soothe_mem.bin'
BASE = 0x180000000

PTRACE_CONT = 7
PTRACE_GETREGS = 12
PTRACE_SETREGS = 13
PTRACE_PEEKUSER = 3
PTRACE_POKEUSER = 6
PTRACE_SINGLESTEP = 9
PTRACE_SEIZE = 0x4206
PTRACE_INTERRUPT = 0x4207
PTRACE_O_TRACECLONE = 2

DR0_OFF = 0x350          # empirically verified debugreg base
DR7_OFF = DR0_OFF + 56   # 0x388
DR7_VAL = 0x000F0001     # L0=1, R/W0=11 (rd/wr), LEN0=11 (4 bytes)

libc = ctypes.CDLL('libc.so.6', use_errno=True)
_mod = open(BIN, 'rb').read()


class UserRegs(ctypes.Structure):
    _fields_ = [(n, ctypes.c_ulonglong) for n in (
        'r15', 'r14', 'r13', 'r12', 'rbp', 'rbx', 'r11', 'r10',
        'r9', 'r8', 'rax', 'rcx', 'rdx', 'rsi', 'rdi', 'orig_rax',
        'rip', 'cs', 'eflags', 'rsp', 'ss', 'fs_base', 'gs_base',
        'ds', 'es', 'fs', 'gs')]


def pt(req, pid, addr=0, data=0):
    libc.ptrace.restype = ctypes.c_long
    r = libc.ptrace(req, pid, ctypes.c_void_p(addr), ctypes.c_void_p(data))
    return None if r == -1 else r


def find_host():
    for p in glob.glob('/proc/[0-9]*'):
        pid = int(os.path.basename(p))
        try:
            cmd = open(f'/proc/{pid}/cmdline', 'rb').read().replace(b'\0', b' ').decode('utf8', 'replace')
            maps = open(f'/proc/{pid}/maps').read()
        except Exception:
            continue
        if 'soothe2' in maps and 'reaper' not in cmd:
            return pid
    return None


def main():
    subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; "
                   "rm -rf /run/user/1000/yabridge-soothe2_x64-*; sleep 1", shell=True)
    proc = subprocess.Popen(['/usr/bin/reaper', '-nosplash', '-ignoreerrors',
                             '-renderproject', '/home/m/soothe-bt/dual_b1q_0.5.rpp'],
                            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    host = None
    t0 = time.time()
    while time.time() - t0 < 30 and not host:
        host = find_host()
        time.sleep(0.002)
    if not host:
        print('NO HOST')
        return 1

    fd = os.open(f'/proc/{host}/mem', os.O_RDONLY)

    def rd(a, n):
        try:
            return os.pread(fd, n, a)
        except OSError:
            return None

    # find ctx & FIR ptr while running (poll till instance exists)
    ctx = firptr = None
    vt = struct.pack('<Q', 0x1824AC210)
    while time.time() - t0 < 25 and firptr is None:
        for line in open(f'/proc/{host}/maps'):
            parts = line.split()
            if 'rw' not in parts[1]:
                continue
            lo, hi = (int(x, 16) for x in parts[0].split('-'))
            CH = 16 * 1024 * 1024
            a = lo
            while a < hi and firptr is None:
                d = rd(a, min(CH + 4096, hi - a))
                if not d:
                    break
                j = d.find(vt)
                while j >= 0:
                    cand = a + j
                    sb = rd(cand + 0x540870, 4)
                    if sb and struct.unpack('<f', sb)[0] > 100:
                        ctx = cand
                        qb = rd(cand + 0x540668, 8)
                        firptr = struct.unpack('<Q', qb)[0]
                        break
                    j = d.find(vt, j + 1)
                a += CH
    if firptr is None:
        print('NO CTX/FIR')
        return 1
    print('host %d ctx %#x fir %#x at %.2fs' % (host, ctx, firptr, time.time() - t0), flush=True)

    seized = {host}
    armed = set()
    pt(PTRACE_SEIZE, host, 0, PTRACE_O_TRACECLONE)
    pt(PTRACE_INTERRUPT, host)
    for _ in range(60):
        try:
            pid, st = os.waitpid(host, os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == host:
            break
        time.sleep(0.001)
    pt(PTRACE_POKEUSER, host, DR0_OFF, 0x18052b8bb)
    pt(PTRACE_POKEUSER, host, DR7_OFF, DR7_VAL)
    armed.add(host)
    print('placeholder armed', flush=True)

    def arm_dbg(tid):
        pass_off=(0x350,)
        for off in pass_off:
            r = pt(PTRACE_POKEUSER, tid, off, 0xdeadbeef if off == DR0_OFF else DR7_VAL)
            v = pt(PTRACE_PEEKUSER, tid, off)
            print('   dbg poke %#x -> r=%s read=%#x' % (off, 'ok' if r is not None else 'EIO', v or 0), flush=True)

    def arm(tid):
        pt(PTRACE_POKEUSER, tid, DR0_OFF, firptr + 43 * 8)   # bin43 re (hot word)
        pt(PTRACE_POKEUSER, tid, DR7_OFF, DR7_VAL)
        v = pt(PTRACE_PEEKUSER, tid, DR7_OFF)
        # bit10 of DR7 reads as always-1 (RA1)
        return v is not None and (v & ~0x400) == (DR7_VAL & ~0x400)

    # stop host briefly to verify arming works at all
    pt(PTRACE_INTERRUPT, host)
    for _ in range(50):
        try:
            pid, st = os.waitpid(host, os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == host:
            break
        time.sleep(0.002)
    arm_dbg(host)
    ok = arm(host)
    print('arm check (stopped):', ok, flush=True)
    pt(PTRACE_CONT, host, 0, 0)

    if ok:
        armed.add(host)

    hits = []
    t_end = time.time() + DUR
    last_sweep = 0.0
    while time.time() < t_end:
        now = time.time()
        if now - last_sweep > 0.03:
            last_sweep = now
            for tid_s in glob.glob(f'/proc/{host}/task/*'):
                tid = int(os.path.basename(tid_s))
                if tid not in seized:
                    if pt(PTRACE_SEIZE, tid, 0, 0) is not None:
                        seized.add(tid)
                    if tid not in armed:
                        pt(PTRACE_INTERRUPT, tid)
                        for _ in range(40):
                            try:
                                pid, st = os.waitpid(tid, os.WUNTRACED | os.WNOHANG)
                            except ChildProcessError:
                                break
                            if pid == tid:
                                break
                            time.sleep(0.001)
                        if arm(tid):
                            armed.add(tid)
                        try:
                            pt(PTRACE_CONT, tid, 0, 0)
                        except OSError:
                            pass
        try:
            pid, status = os.waitpid(-1, os.WSTOPPED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == 0:
            time.sleep(0.0005)
            continue
        sig = status >> 8
        if os.WIFEXITED(status) or os.WIFSIGNALED(status):
            continue
        if os.WIFSTOPPED(pid) and sig == signal.SIGTRAP:
            regs = UserRegs()
            if pt(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs)) is None:
                continue
            rec = dict(rip=regs.rip, rsp=regs.rsp, tid=pid,
                       ret=struct.unpack('<Q', rd(regs.rsp, 8) or b'\0' * 8)[0])
            hits.append(rec)
            if len(hits) <= 15:
                print('WATCH HIT rip=%#x ret=%#x' % (rec['rip'], rec['ret']), flush=True)
            # pass trap: clear DR6 (poke 0xffff0ff0), single-step, continue
            pt(PTRACE_POKEUSER, pid, DR0_OFF + 48, 0xFFFF0FF0)  # reset DR6
            regs.eflags |= 0x100
            pt(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs))
            pt(PTRACE_SINGLESTEP, pid, 0, 0)
            try:
                os.waitpid(pid, os.WUNTRACED)
            except ChildProcessError:
                pass
            pt(PTRACE_CONT, pid, 0, 0)
        elif os.WIFSTOPPED(pid):
            pt(PTRACE_CONT, pid, 0, sig if 0 < sig < 32 else 0)

    print('total watch hits:', len(hits))
    import json
    json.dump(hits, open('/tmp/opencode/fnwatch/hits.json', 'w'), indent=1, default=str)
    proc.kill()
    return 0


if __name__ == '__main__':
    os.makedirs('/tmp/opencode/fnwatch', exist_ok=True)
    sys.exit(main())