#!/usr/bin/env python3
"""fnhw.py — hardware-breakpoint tracer (DR0) for the yabridge host.
Safe for un-traced threads (code is never patched). Auto-arms new threads.

Env: FN_ADDR (target VA), FN_DUR (seconds), FN_SKIP (skip first N hits)
"""
import ctypes
import glob
import json
import os
import signal
import struct
import subprocess
import sys
import time

FN = int(os.environ.get('FN_ADDR', '0x180529FE0'), 16)
DUR = float(os.environ.get('FN_DUR', '15'))

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

libc = ctypes.CDLL('libc.so.6', use_errno=True)


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))
    if r == -1:
        return None, ctypes.get_errno()
    return r, 0


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 find_dr0_off(pid):
    """Locate DR0 slot inside the USER area by probing."""
    probe = 0x1234567890abcdef
    for off in range(0x200, 0x400, 8):
        r, e = pt(PTRACE_POKEUSER, pid, off, probe)
        if r is None:
            continue
        v, _ = pt(PTRACE_PEEKUSER, pid, off)
        if v == (probe & 0xFFFFFFFFFFFFFFFF):
            # restore zero & sanity-check neighbours exist
            pt(PTRACE_POKEUSER, pid, off, 0)
            return off
    return None


def arm(tid, dr0_off):
    """Arm DR0/DR7 - requires tid to be STOPPED; caller handles stop/cont."""
    pt(PTRACE_POKEUSER, tid, dr0_off, FN)
    # DR7 (index 7): L0=1, RW0=00 (exec), LEN0=00
    pt(PTRACE_POKEUSER, tid, dr0_off + 7 * 8, 0x1)
    v, _ = pt(PTRACE_PEEKUSER, tid, dr0_off + 7 * 8)
    return v == 1


def disarm(tid, dr0_off):
    pt(PTRACE_POKEUSER, tid, dr0_off + 7 * 8, 0)
    pt(PTRACE_POKEUSER, tid, dr0_off, 0)


def arm_stopped(tid, dr0_off):
    """INTERRUPT -> (bounded) wait -> arm -> CONT. True if DR7 verified."""
    pt(PTRACE_INTERRUPT, tid)
    for _ in range(30):
        try:
            pid, st = os.waitpid(tid, os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            return False
        if pid == tid:
            break
        time.sleep(0.002)
    r0, e0 = pt(PTRACE_POKEUSER, tid, dr0_off, FN)
    r7, e7 = pt(PTRACE_POKEUSER, tid, dr0_off + 7 * 8, 0x1)
    v, _ = pt(PTRACE_PEEKUSER, tid, dr0_off + 7 * 8)
    print('   arm tid=%d poke_dr0=%s(e%s) poke_dr7=%s(e%s) dr7_read=%s'
          % (tid, 'ok' if r0 is not None else 'FAIL', e0,
             'ok' if r7 is not None else 'FAIL', e7, hex(v or 0)), flush=True)
    pt(PTRACE_CONT, tid, 0, 0)
    return v == 1


def main():
    subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; 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
    print('host %d at %.3fs' % (host, time.time() - t0), flush=True)

    seized = []
    for tid in [int(os.path.basename(p)) for p in glob.glob(f'/proc/{host}/task/*')] or [host]:
        if tid not in seized:
            r, e = pt(PTRACE_SEIZE, tid, 0, 0)
            if r is not None or e != 3:
                seized.append(tid)

    # stop one thread briefly to locate DR0 offset
    tgt = seized[0]
    pt(PTRACE_INTERRUPT, tgt)
    os.waitpid(tgt, os.WUNTRACED)
    dr0_off = find_dr0_off(tgt)
    print('DR0 user-offset:', hex(dr0_off) if dr0_off else 'NOT FOUND', flush=True)
    if not dr0_off:
        return 1
    pt(PTRACE_CONT, tgt, 0, 0)
    n_ok = 0
    for tid in seized:
        if arm_stopped(tid, dr0_off):
            n_ok += 1
    print('armed %#x on %d/%d tids (DR7 verified)' % (FN, n_ok, len(seized)), flush=True)

    log = []
    hits = 0
    skip = int(os.environ.get('FN_SKIP', '0'))
    t_end = time.time() + DUR
    last_resweep = 0.0
    while time.time() < t_end:
        now = time.time()
        if now - last_resweep > 0.03:
            last_resweep = now
            for tid_s in glob.glob(f'/proc/{host}/task/*'):
                tid = int(os.path.basename(tid_s))
                if tid not in seized:
                    r, e = pt(PTRACE_SEIZE, tid, 0, 0)
                    if r is not None or e != 3:
                        seized.append(tid)
                        if arm_stopped(tid, dr0_off):
                            print('+tid', tid, flush=True)
        try:
            pid, status = os.waitpid(-1, os.WSTOPPED | os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == 0:
            time.sleep(0.0005)
            continue
        sig = status >> 8
        ev = status >> 16
        if os.WIFEXITED(status) or os.WIFSIGNALED(status):
            continue
        if ev == 3 or ev == 1:
            pt(PTRACE_CONT, pid, 0, 0)
            continue
        if os.WIFSTOPPED(pid) and sig == signal.SIGTRAP:
            regs = UserRegs()
            r, _ = pt(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs))
            if r is None:
                continue
            if regs.rip == FN or regs.rip == FN + 1:
                rip = FN
                ret, _ = pt(PTRACE_PEEKDATA, pid, regs.rsp)
                rec = dict(ctx=regs.rcx, a2=regs.rdx, cnt=regs.r8 & 0xffffffff,
                           r9=regs.r9 & 0xffffffff, ret=ret,
                           rbx=regs.rbx, r12=regs.r12, r13=regs.r13,
                           r14=regs.r14, r15=regs.r15, rsp=regs.rsp, tid=pid)
                if hits >= skip:
                    log.append(rec)
                    print('HIT ctx=%#x a2=%#x cnt=%#x r9d=%#x ret=%#x'
                          % (rec['ctx'], rec['a2'], rec['cnt'], rec['r9'], rec['ret']),
                          flush=True)
                hits += 1
                # pass bp: clear DR0 temporarily, single-step, re-arm
                pt(PTRACE_POKEUSER, pid, dr0_off + 7 * 8, 0)
                regs.eflags |= 0x100  # TF
                pt(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs))
                pt(PTRACE_SINGLESTEP, pid, 0, 0)
                os.waitpid(pid, os.WUNTRACED)
                regs2 = UserRegs()
                pt(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs2))
                regs2.eflags &= ~0x100
                pt(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs2))
                arm(pid, dr0_off)
                pt(PTRACE_CONT, pid, 0, 0)
            else:
                # stray SIGTRAP (wine internal): forward
                pt(PTRACE_CONT, pid, 0, signal.SIGTRAP)
        elif os.WIFSTOPPED(pid):
            pt(PTRACE_CONT, pid, 0, sig if 0 < sig < 32 else 0)

    print('total hits:', hits, 'logged:', len(log))
    json.dump(log, open('/tmp/opencode/fntrace/hwhits.json', 'w'), indent=1)
    for tid in seized:
        disarm(tid, dr0_off)
    proc.kill()
    return 0


if __name__ == '__main__':
    sys.exit(main())