#!/usr/bin/env python3
"""fntrace3.py — decisive INT3 experiment: seize ALL tids within milliseconds
of host spawn, arm breakpoint, log EVERY waitpid event for N seconds."""
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', '25'))

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_LISTEN = 0x4208
PTRACE_O_TRACECLONE = 2

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 main():
    rpp = '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    subprocess.run('pkill -9 -x reaper; pkill -9 -f \'[y]abridge\'; sleep 1', shell=True)
    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
    print('host %d at %.3fs' % (host, time.time() - t0), flush=True)

    seized = []
    for tid in [host] + [int(os.path.basename(p)) for p in glob.glob(f'/proc/{host}/task/*')]:
        if tid in seized:
            continue
        r, e = pt(PTRACE_SEIZE, tid, 0, PTRACE_O_TRACECLONE)
        if r is None and e == 3:  # ESRCH - gone
            continue
        seized.append(tid)
        # interrupt to allow poke below (poke needs SOME stopped thread)
    print('seized:', seized, flush=True)
    # stop one thread to enable POKETEXT
    tgt = seized[0]
    pt(PTRACE_INTERRUPT, tgt)
    os.waitpid(tgt, os.WUNTRACED)
    orig, _ = pt(PTRACE_PEEKDATA, tgt, FN)
    cc = (orig & ~0xFF) | 0xCC
    pt(PTRACE_POKETEXT, tgt, FN, cc)
    print('armed orig=%#x' % orig, flush=True)
    # seize any tids spawned meanwhile
    for tid in [int(os.path.basename(p)) for p in glob.glob(f'/proc/{host}/task/*')]:
        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)
    for tid in seized:
        try:
            pt(PTRACE_CONT, tid, 0, 0)
        except Exception:
            pass

    log = []
    events = []
    hits = 0
    t_end = time.time() + DUR
    last_resweep = 0.0
    while time.time() < t_end:
        now = time.time()
        if now - last_resweep > 0.05:
            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, PTRACE_O_TRACECLONE)
                    if r is not None or e != 3:
                        seized.append(tid)
                        print('+tid', tid, flush=True)
        pid, status = os.waitpid(-1, os.WSTOPPED | os.WUNTRACED | os.WNOHANG)
        if pid == 0:
            time.sleep(0.001)
            continue
        sig = status >> 8
        ev = status >> 16
        events.append((time.time() - t0, pid, hex(status), ev))
        if pid not in seized:
            r, e = pt(PTRACE_SEIZE, pid, 0, PTRACE_O_TRACECLONE)
            if r is not None or e != 3:
                seized.append(pid)
        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 - 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)
                log.append(rec)
                hits += 1
                print('HIT ctx=%#x a2=%#x cnt=%#x ret=%#x'
                      % (regs.rcx, regs.rdx, regs.r8 & 0xffffffff, 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.WIFEXITED(status) or os.WIFSIGNALED(status):
            continue
        else:
            pt(PTRACE_CONT, pid, 0, sig if 0 < sig < 32 else 0)

    print('total hits:', hits, 'events:', len(events))
    json.dump(dict(log=log, events=events[:400]),
              open('/tmp/opencode/fntrace/hits3.json', 'w'), indent=1)
    proc.kill()
    return 0


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