#!/usr/bin/env python3
"""fnall.py — safe full-thread INT3 tracer.

Invariant: every thread is SEIZE+INTERRUPT-stopped BEFORE the breakpoint is
armed, so any later clone descends from a traced thread and its SIGTRAPs come
to us instead of killing the host.

Env: FN_ADDR, FN_DUR, FN_MAXHITS
"""
import ctypes
import glob
import json
import os
import signal
import struct
import numpy as np
import subprocess
import sys
import time

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

PTRACE_CONT = 7
PTRACE_GETREGS = 12
PTRACE_SETREGS = 13
PTRACE_PEEKDATA = 2
PTRACE_POKETEXT = 4
PTRACE_SINGLESTEP = 9
PTRACE_SEIZE = 0x4206
PTRACE_INTERRUPT = 0x4207
PTRACE_O_TRACECLONE = 0x00000002

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():
    # fastest signal: any non-reaper child of a reaper process
    reapers = []
    for p in glob.glob('/proc/[0-9]*'):
        try:
            cmd = open(f'/proc/{p}/cmdline', 'rb').read()
            if b'reaper' in cmd and b'yabridge' not in cmd:
                reapers.append(int(os.path.basename(p)))
        except Exception:
            continue
    for p in glob.glob('/proc/[0-9]*'):
        pid = int(os.path.basename(p))
        try:
            st = open(f'/proc/{pid}/stat').read().split()
            ppid = int(st[3])
        except Exception:
            continue
        if ppid in reapers and pid not in reapers:
            return pid
    return None


def drain_waits():
    out = []
    while True:
        try:
            pid, status = os.waitpid(-1, os.WSTOPPED | os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            break
        if pid == 0:
            break
        out.append((pid, status))
    return out


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)

    # phase 1: seize main, interrupt, wait stop
    seized = []
    r, e = pt(PTRACE_SEIZE, host, 0, PTRACE_O_TRACECLONE)
    if r is None:
        print('seize fail', e)
        return 1
    seized.append(host)
    pt(PTRACE_INTERRUPT, host)

    # phase 2: pump events, seize newcomers until set stabilizes
    stable_until = time.time() + 2.5
    deadline = time.time() + 8.0
    while time.time() < deadline:
        for pid, st in drain_waits():
            pass  # they stay stopped; we hold them
        grew = False
        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, PTRACE_O_TRACECLONE)
                if r is not None or e != 3:
                    seized.append(tid)
                    pt(PTRACE_INTERRUPT, tid)
                    grew = True
        if grew:
            stable_until = time.time() + 0.5
        elif time.time() > stable_until:
            break
        time.sleep(0.002)
    print('stopped %d threads' % len(seized), flush=True)

    orig, _ = pt(PTRACE_PEEKDATA, host, FN)
    cc = (orig & ~0xFF) | 0xCC
    r, e = pt(PTRACE_POKETEXT, host, FN, cc)
    if r is None:
        print('arm fail', e)
        return 1
    print('armed %#x' % FN, flush=True)

    log = []
    hits = 0
    EVENTS = []
    for tid in seized:
        pt(PTRACE_CONT, tid, 0, 0)

    t_end = time.time() + DUR
    last_resweep = time.time()
    while hits < MAXHITS and time.time() < t_end:
        try:
            pid, status = os.waitpid(-1, os.WSTOPPED | os.WUNTRACED | os.WNOHANG)
        except ChildProcessError:
            break
        EVENTS.append((time.time()-t0, pid, hex(status), status >> 16))
        if pid == 0:
            if time.time() - last_resweep > 0.05:
                last_resweep = time.time()
                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, PTRACE_O_TRACECLONE)
                        if r is not None or e != 3:
                            seized.append(tid)
                            pt(PTRACE_INTERRUPT, tid)
                            drain_waits()
                            pt(PTRACE_CONT, tid, 0, 0)
            time.sleep(0.0005)
            continue
        sig = status >> 8
        ev = status >> 16
        if os.WIFEXITED(status) or os.WIFSIGNALED(status):
            EVENTS.append((time.time()-t0, pid, hex(status), ev))
            continue
        if ev == 3 or ev == 1:
            pt(PTRACE_CONT, pid, 0, 0)
            continue
        if os.WIFSTOPPED(pid) and sig == signal.SIGTRAP and ev == 0:
            regs = UserRegs()
            r, _ = pt(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs))
            if r is None:
                continue
            if regs.rip - 1 == 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,
                           dump_in=[], dump_out=[], dump_r14=[], dump_rbx=[])
                # dump arrays referenced by registers (process is stopped)
                fdm = os.open(f'/proc/{host}/mem', os.O_RDONLY)

                def rdarr(ptr, n=2049):
                    try:
                        b = os.pread(fdm, n * 4, ptr)
                        return np.frombuffer(b, dtype='<f4').astype(np.float32).tolist()
                    except OSError:
                        return []
                rec['dump_in'] = rdarr(regs.rcx)
                rec['dump_out'] = rdarr(regs.rdx)
                if FN != 0x18052dbc0:
                    rec['dump_r14'] = rdarr(regs.r14)
                os.close(fdm)
                log.append(rec)
                hits += 1
                if hits <= 20:
                    print('HIT ctx=%#x a2=%#x cnt=%#x ret=%#x'
                          % (rec['ctx'], rec['a2'], rec['cnt'], rec['ret']), flush=True)
                pt(PTRACE_POKETEXT, pid, FN, orig)
                regs.rip = FN
                pt(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs))
                pt(PTRACE_SINGLESTEP, pid, 0, 0)
                os.waitpid(pid, os.WUNTRACED)
                pt(PTRACE_POKETEXT, pid, FN, cc)
                pt(PTRACE_CONT, pid, 0, 0)
            else:
                pt(PTRACE_CONT, pid, 0, 0)
        elif os.WIFSTOPPED(pid):
            pt(PTRACE_CONT, pid, 0, sig if 0 < sig < 32 else 0)

    pt(PTRACE_POKETEXT, host, FN, orig)
    print('total hits:', hits)
    json.dump(dict(log=log, events=EVENTS[:8000], nseized=len(seized)),
              open('/tmp/opencode/fntrace/allhits.json', 'w'), indent=1)
    proc.kill()
    return 0


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