#!/usr/bin/env python3
"""fntrace.py — ptrace INT3 tracer for FUN_180529fe0 (mask chain vtbl slot 6)
in the live yabridge host. The wine module is mapped at its preferred base,
so dump VAs == runtime addresses (verified: exec map 0x180001000-0x181baa000).

On each hit logs: RIP, RCX (ctx), RDX, R8D (count), R9D, [RSP] (return addr),
plus xmm0/xmm1 low scalars if available via GETFPREGS (skipped: not portable).

Usage: python3 scripts/fntrace.py [rpp] [nhits]
"""
import ctypes
import glob
import os
import signal
import struct
import subprocess
import sys
import time

FN = int(os.environ.get('FN_ADDR', '0x180529FE0'), 16)
SNAPDIR = '/tmp/opencode/fntrace'

PTRACE_TRACEME = 0
PTRACE_PEEKTEXT = 1
PTRACE_PEEKDATA = 2
PTRACE_POKETEXT = 4
PTRACE_CONT = 7
PTRACE_SINGLESTEP = 9
PTRACE_GETREGS = 12
PTRACE_SETREGS = 13
PTRACE_ATTACH = 16
PTRACE_DETACH = 17
PTRACE_SEIZE = 0x4206
PTRACE_INTERRUPT = 0x4207
PTRACE_O_TRACECLONE = 0x00000002
PTRACE_EVENT_CLONE = 3  # status >> 16 == 4 (event+1)? actually event = status>>16, CLONE==3 -> 4? use raw compare below
PTRACE_EVENT_FORK = 1


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')]


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


def ptrace(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:
        e = ctypes.get_errno()
        if req not in (PTRACE_PEEKTEXT, PTRACE_PEEKDATA):
            raise OSError(e, f'ptrace({req:#x},{pid}) failed')
        return None
    return 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():
    rpp = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    nhits = int(sys.argv[2]) if len(sys.argv) > 2 else 24
    render = '--render' in sys.argv
    os.makedirs(SNAPDIR, exist_ok=True)
    subprocess.run('pkill -9 -x reaser 2>/dev/null; pkill -9 -x reaper 2>/dev/null; '
                   "pkill -9 -f '[y]abridge' 2>/dev/null; sleep 1", shell=True)
    if render:
        proc = subprocess.Popen(
            ['/usr/bin/reaper', '-nosplash', '-ignoreerrors',
             '-renderproject', rpp],
            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    else:
        proc = subprocess.Popen(
            ['/usr/bin/reaper', '-nosplash', '-ignoreerrors', rpp,
             '/home/m/re-tools/play_loop.lua'],
            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    t0 = time.time()
    host = None
    while time.time() - t0 < 60 and not host:
        host = find_host()
        time.sleep(0.05)
    if not host:
        print('NO HOST')
        return 1
    print('host', host, 'at %.1fs' % (time.time() - t0), flush=True)

    tids = [int(os.path.basename(p)) for p in glob.glob(f'/proc/{host}/task/*')]
    print('tids:', tids, flush=True)

    seized = []
    for tid in tids:
        try:
            ptrace(PTRACE_SEIZE, tid, 0, PTRACE_O_TRACECLONE)
            seized.append(tid)
        except OSError as e:
            print('seize fail', tid, e)
    if not seized:
        return 1
    # stop everyone
    stopped = []
    for tid in seized:
        try:
            ptrace(PTRACE_INTERRUPT, tid)
            os.waitpid(tid, os.WUNTRACED)
            stopped.append(tid)
        except (OSError, ChildProcessError):
            pass

    orig = ptrace(PTRACE_PEEKTEXT, stopped[0], FN)
    cc = (orig & ~0xFF) | 0xCC
    ptrace(PTRACE_POKETEXT, stopped[0], FN, cc)
    print('breakpoint armed at %#x (orig=%#x)' % (FN, orig), flush=True)

    for tid in stopped:
        try:
            ptrace(PTRACE_CONT, tid, 0, 0)
        except OSError:
            pass

    hits = 0
    log = []
    import select
    import time as _t
    t_last = _t.time()
    idle_deadline = float(os.environ.get('FNTRACE_IDLE', '20'))
    while hits < nhits and _t.time() - t_last < idle_deadline:
        try:
            pid, status = os.waitpid(-1, os.WUNTRACED | os.WSTOPPED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == 0:
            _t.sleep(0.005)
            continue
        sig = status >> 8
        if os.WIFEXITED(pid and status or status) or os.WIFSIGNALED(status):
            # thread/process exited (normal in wine): forget it
            if pid in seized:
                seized.remove(pid)
            if pid in stopped:
                stopped.remove(pid)
            continue
        if not os.WIFSTOPPED(pid):
            continue
        t_last = _t.time()
        ev = status >> 16
        if ev == PTRACE_EVENT_CLONE or ev == PTRACE_EVENT_FORK:
            if pid not in seized:
                seized.append(pid)
            ptrace(PTRACE_CONT, pid, 0, 0)
            continue
        if sig == signal.SIGTRAP:
            regs = UserRegs()
            try:
                ptrace(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs))
            except OSError:
                ptrace(PTRACE_CONT, pid, 0, 0)
                continue
            if regs.rip - 1 == FN:
                ret = ptrace(PTRACE_PEEKDATA, pid, regs.rsp)
                rec = dict(rip=regs.rip - 1, ctx=regs.rcx, a2=regs.rdx,
                           cnt=regs.r8 & 0xffffffff, r9=regs.r9 & 0xffffffff,
                           ret=ret, tid=pid,
                           rbx=regs.rbx, rbp_=regs.rbp, rsi=regs.rsi, rdi=regs.rdi)
                log.append(rec)
                hits += 1
                print(f'hit {hits}: tid={pid} ctx={regs.rcx:#x} '
                      f'a2={regs.rdx:#x} cnt={regs.r8:#x} r9d={regs.r9:#x} '
                      f'ret={ret:#x}', flush=True)
                # step over int3
                lo = struct.unpack('<Q', struct.pack('<Q', orig ^ ((orig ^ cc) & 0xFF)))[0]
                ptrace(PTRACE_POKETEXT, pid, FN, orig)
                regs.rip = FN
                ptrace(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs))
                ptrace(PTRACE_SINGLESTEP, pid, 0, 0)
                os.waitpid(pid, os.WUNTRACED)
                ptrace(PTRACE_POKETEXT, pid, FN, cc)
                ptrace(PTRACE_CONT, pid, 0, 0)
            else:
                ptrace(PTRACE_CONT, pid, 0, 0)
        else:
            # other stop signals: deliver and continue
            ptrace(PTRACE_CONT, pid, 0, sig if 0 < sig < 0x20 else 0)

    # cleanup: remove breakpoint, detach
    print('done hits=', hits, 'cleaning up...', flush=True)
    for tid in seized:
        try:
            ptrace(PTRACE_INTERRUPT, tid)
            os.waitpid(tid, os.WUNTRACED)
        except (OSError, ChildProcessError):
            continue
    try:
        ptrace(PTRACE_POKETEXT, stopped[0], FN, orig)
    except Exception as e:
        print('restore fail', e)
    for tid in seized:
        try:
            ptrace(PTRACE_DETACH, tid, 0, 0)
        except OSError:
            pass
    import json
    json.dump(log, open(f'{SNAPDIR}/hits.json', 'w'), indent=1)
    print('saved', f'{SNAPDIR}/hits.json')
    if proc.poll() is None:
        proc.kill()
    return 0


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