#!/usr/bin/env python3
"""fnwatch4.py — 4 hardware watchpoints simultaneously:
  DR0/DR1: dwords of the FIR pointer slot ctx+0x540668 (catches loader of ptr)
  DR2    : kernel word bin43.re
  DR3    : scratch word bin43
RW=read/write (traps readers AND writers), LEN=4B."""
import ctypes
import glob
import hashlib
import json
import os
import signal
import struct
import subprocess
import sys
import time

DUR = float(os.environ.get('WATCH_DUR', '20'))
PTRACE_CONT = 7
PTRACE_GETREGS = 12
PTRACE_SETREGS = 13
PTRACE_PEEKUSER = 3
PTRACE_POKEUSER = 6
PTRACE_SEIZE = 0x4206
PTRACE_INTERRUPT = 0x4207
PTRACE_O_TRACECLONE = 2

DR_BASE = 0x350          # u_debugreg[0]
DR6_OFF = DR_BASE + 48   # 0x388 - wait: idx6=+48 -> that IS DR6 slot? no:
# u_debugreg[0..7]: DR0@+0, DR1@+8, DR2@+16, DR3@+24, DR4@+32, DR5@+40,
#                   DR6@+48, DR7@+56
DR7_OFF = DR_BASE + 56
DR7_VAL = 0xFFFF000F     # L0-3 enabled, all RW=11, all LEN=11 (4B)

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))
    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)
    rpp = '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    wav = rpp.replace('.rpp', '.wav')
    wt0 = os.path.getmtime(wav) if os.path.exists(wav) else 0
    proc = subprocess.Popen(['/usr/bin/reaper', '-nosplash', '-ignoreerrors',
                             '-renderproject', 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

    ctx = firptr = scrptr = 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
                        firptr = struct.unpack('<Q', rd(cand + 0x540668, 8))[0]
                        scrptr = struct.unpack('<Q', rd(cand + 0x540628, 8))[0] or firptr
                        break
                    j = d.find(vt, j + 1)
                a += CH
    if firptr is None:
        print('NO CTX')
        return 1
    print('host %d ctx %#x fir %#x scr %#x at %.2fs' %
          (host, ctx, firptr, scrptr, time.time() - t0), flush=True)

    seized = {host}
    armed = set()
    pt(PTRACE_SEIZE, host, 0, PTRACE_O_TRACECLONE)

    def arm(tid):
        okall = True
        ovf = struct.unpack('<Q', rd(ctx + 0x5406f8, 8))[0]
        trk = struct.unpack('<Q', rd(ctx + 0x540768, 8))[0]
        rax = struct.unpack('<Q', rd(ctx + 0x5407f8, 8))[0]
        targets = [(DR_BASE + 0, ovf + 8192),
                   (DR_BASE + 8, trk + 43 * 4),
                   (DR_BASE + 16, rax + 43 * 4),
                   (DR_BASE + 24, ctx + 0x2404dc)]
        for off, addr in targets:
            pt(PTRACE_POKEUSER, tid, off, addr)
            v = pt(PTRACE_PEEKUSER, tid, off)
            if v != addr:
                okall = False
        pt(PTRACE_POKEUSER, tid, DR7_OFF, DR7_VAL)
        return okall

    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)
    ok = arm(host)
    print('arm(host)=%s' % 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, PTRACE_O_TRACECLONE) is not None:
                        seized.add(tid)
                    pt(PTRACE_INTERRUPT, tid)
                    for _ in range(40):
                        try:
                            wpid, _st = os.waitpid(tid, os.WUNTRACED | os.WNOHANG)
                        except ChildProcessError:
                            break
                        if wpid == 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
            dr6v = pt(PTRACE_PEEKUSER, pid, DR_BASE + 48)
            rec = dict(rip=regs.rip, rsp=regs.rsp, tid=pid,
                       dr6=dr6v,
                       ret=struct.unpack('<Q', rd(regs.rsp, 8) or b'\0' * 8)[0])
            hits.append(rec)
            if len(hits) <= 20:
                print('HIT rip=%#x ret=%#x dr6=%#x'
                      % (rec['rip'], rec['ret'], dr6v or 0), flush=True)
            # clear DR6 and single-step past
            pt(PTRACE_POKEUSER, pid, DR_BASE + 48, 0xFFFF0FF0)
            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)

        if not hasattr(main, '_wv') and os.path.exists(wav) and os.path.getmtime(wav) > wt0:
            main._wv = True
            print('wav fresh at %.2fs' % (time.time() - t0), flush=True)
    if not hasattr(main, '_wv'):
        main._wv = False
    fresh = main._wv
    print('hits=%d render_fresh=%s armed=%d/%d' % (len(hits), fresh, len(armed), len(seized)))
    json.dump(hits[:200], open('/tmp/opencode/fnwatch/hits4.json', 'w'),
              indent=1, default=str)
    proc.kill()
    return 0


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