#!/usr/bin/env python3
"""iat_name.py — resolve imported-function names for bigkernel dispatch targets.
Reads runtime IAT values, finds owning module, parses PE exports."""
import struct, subprocess, sys, time
import glob, os

BASE=0x180000000
data=open('/home/m/re-tools/soothe_mem.bin','rb').read()

def rd(fd,a,n):
    try: return os.pread(fd,n,a)
    except OSError: return None

def find_host():
    for p in glob.glob('/proc/[0-9]*'):
        pid=int(os.path.basename(p))
        try:
            cmd=open('/proc/%d/cmdline'%pid,'rb').read().replace(b'\0',b' ').decode('utf8','replace')
            maps=open('/proc/%d/maps'%pid).read()
        except Exception: continue
        if 'soothe2' in maps and 'reaper' not in cmd: return pid
    return None

def iat_slot(stub):
    # pattern: mov rax,[rip+rel] (48 8b 05 rel32)
    off=stub-BASE
    b=data[off:off+7]
    if b[:2]!=b'\x48\x8b': return None
    rel=struct.unpack('<i',b[2:6])[0]
    return stub+6+rel

def pe_exports(path):
    """Parse PE export table -> {name: rva}"""
    try:
        f=open(path,'rb').read()
    except Exception:
        return {}
    if f[:2]!=b'MZ': return {}
    pe=struct.unpack('<I',f[0x3c:0x40])[0]
    if f[pe:pe+4]!=b'PE\0\0': return {}
    nsec=struct.unpack('<H',f[pe+6:pe+8])[0]
    optsz=struct.unpack('<H',f[pe+20:pe+22])[0]
    magic=struct.unpack('<H',f[pe+24:pe+26])[0]
    ddir=pe+24+(0x70 if magic==0x20b else 0x60)+0*8   # data dir[0]=export
    exp_rva,exp_sz=struct.unpack('<II',f[ddir:ddir+8])
    if not exp_rva: return {}
    # sections
    secs=[]
    so=pe+24+optsz
    for i in range(nsec):
        s=f[so+i*40:so+i*40+40]
        va,sz=struct.unpack('<II',s[12:20])
        raw,rsz=struct.unpack('<II',s[20:28])
        secs.append((va,sz,raw,rsz))
    def r2o(rva):
        for va,sz,raw,rsz in secs:
            if va<=rva<va+max(sz,rsz): return raw+(rva-va)
        return None
    eo=r2o(exp_rva)
    if eo is None: return {}
    nnames=struct.unpack('<I',f[eo+24:eo+28])[0]
    nrva=struct.unpack('<I',f[eo+32:eo+36])[0]
    names_rva=struct.unpack('<I',f[eo+32+4:eo+32+8])[0]
    funcs_rva=struct.unpack('<I',f[eo+28:eo+32])[0]
    no=r2o(names_rva); fo=r2o(funcs_rva)
    out={}
    if no is None or fo is None: return {}
    for i in range(nnames):
        nrva_i=struct.unpack('<I',f[no+i*4:no+i*4+4])[0]
        noff=r2o(nrva_i)
        if noff is None: continue
        end=f.find(b'\0',noff)
        nm=f[noff:end].decode('ascii','replace')
        frva=struct.unpack('<I',f[fo+i*4:fo+i*4+4])[0]
        out[nm]=frva
    return out

subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; "
               "rm -rf /run/user/1000/yabridge-soothe2_x64-*; sleep 1", shell=True)
proc=subprocess.Popen(['/usr/bin/reaper','-nosplash','-ignoreerrors','-renderproject','/tmp/opencode/multi.rpp'],
                      stdout=subprocess.DEVNULL,stderr=subprocess.STDOUT)
t0=time.time(); host=None
while time.time()-t0<30 and not host:
    host=find_host(); time.sleep(0.002)
print('host',host,flush=True)
import signal as sg
os.kill(host,sg.SIGSTOP)
fd=os.open('/proc/%d/mem'%host,os.O_RDONLY)

# build module map
mods=[]
for line in open('/proc/%d/maps'%host):
    parts=line.split()
    if len(parts)<6 or 'x' not in parts[1]: continue
    lo,hi=(int(x,16) for x in parts[0].split('-'))
    mods.append((lo,hi,parts[5]))
print('modules:',len(mods))

def owner(addr):
    for lo,hi,path in mods:
        if lo<=addr<hi: return (lo,addr-lo,path)
    return None

targets={}
TBL={0x180140b30:0x1826176c8,0x180140b60:0x182617708,0x1801409e0:0x182617508,
     0x180140ad0:0x182617648,0x180140a40:0x182617588}
for stub,tbl in TBL.items():
    v=rd(fd,tbl+32,8)
    tgt=struct.unpack('<Q',v)[0] if v else 0
    # second level: tgt code = mov rax,[rip+rel]; jmp rax -> IAT slot
    print('  L2: tgt=%x' % tgt, flush=True)
    if 0x180000000 <= tgt < 0x187000000:
        off2=tgt-BASE
        b2=data[off2:off2+7] if 0<=off2<len(data)-7 else rd(tgt,7)
        if b2[:2]==b'\x48\x8b':
            rel2=struct.unpack('<i',b2[3:7])[0]
            slot=tgt+7+rel2
            print('  L2: b2=%s rel2=%x slot=%x' % (b2.hex(),rel2&0xffffffff,slot), flush=True)
            fv=rd(fd,slot,8)
            if fv:
                tgt=struct.unpack('<Q',fv)[0]
            else:
                print('   slot read FAIL',slot,flush=True)
        else:
            print('   no mov-rax pattern at %x: %s'%(tgt,b2[:3].hex() if b2 else '-'),flush=True)
    ow=owner(tgt)
    print('%x idx4->%x runtime=%x owner=%s' % (stub,tbl,tgt,ow[2] if ow else '?'),flush=True)
    if ow:
        lo,rva,path=ow
        exps=pe_exports(path)
        best=[nm for nm,r in exps.items() if r==rva]
        print('    export:',best,flush=True)
os.close(fd)