"""Read-only Go pclntab and direct-call inventory for an explicitly selected PE.

This recovers compiled function names and machine-code call edges, not source
code or proof that a handler succeeds. CmdIDs require separate wire evidence.
"""
import argparse
import hashlib
import json
import pathlib
import struct

from capstone import Cs, CS_ARCH_X86, CS_MODE_64
from capstone.x86 import X86_OP_IMM


def dump(path, output):
    data = path.read_bytes()
    pe = struct.unpack_from('<I', data, 0x3c)[0]
    if data[:2] != b'MZ' or data[pe:pe+4] != b'PE\0\0':
        raise ValueError('Expected PE image')
    count = struct.unpack_from('<H', data, pe+6)[0]
    optional = pe+24
    if struct.unpack_from('<H', data, optional)[0] != 0x20b:
        raise ValueError('Expected PE32+')
    image_base = struct.unpack_from('<Q', data, optional+24)[0]
    sections = []
    section_table = optional+struct.unpack_from('<H', data, pe+20)[0]
    for i in range(count):
        pos = section_table+i*40
        name = data[pos:pos+8].rstrip(b'\0').decode('ascii')
        vsize, rva, size, raw = struct.unpack_from('<IIII', data, pos+8)
        sections.append((name, rva, size, raw))
    def offset(va):
        for _, rva, size, raw in sections:
            if rva <= va-image_base < rva+size:
                return raw+va-image_base-rva
        raise ValueError(f'Unmapped VA {va:#x}')
    candidates = []
    start = 0
    while (start := data.find(b'\xf1\xff\xff\xff\0\0\x01\x08', start)) >= 0:
        try:
            nfunc, nfiles, text, names, cu, files, pctab, pcln = struct.unpack_from('<8Q', data, start+8)
            if 0 < nfunc < 1000000 and 0 < nfiles < 1000000 and start+pcln+(nfunc+1)*8 < len(data):
                if text == 0:
                    text = image_base+next(rva for name,rva,_,_ in sections if name == '.text')
                offset(text)
                candidates.append((start, nfunc, text, names, pcln))
        except (ValueError, struct.error):
            pass
        start += 8
    if len(candidates) != 1:
        raise ValueError(f'Expected one valid Go1.20+ pclntab, found {len(candidates)}')
    header, nfunc, text, names, pcln = candidates[0]
    funcs = []
    for i in range(nfunc):
        entry, info = struct.unpack_from('<II', data, header+pcln+i*8)
        end = struct.unpack_from('<I', data, header+pcln+(i+1)*8)[0]
        nameoff = struct.unpack_from('<i', data, header+pcln+info+4)[0]
        pos = header+names+nameoff
        name = data[pos:data.index(b'\0', pos)].decode('utf8')
        funcs.append({'name': name, 'va': hex(text+entry), 'rva': hex(text+entry-image_base), 'size': end-entry})
    byva = {int(f['va'], 16): f['name'] for f in funcs}
    cs = Cs(CS_ARCH_X86, CS_MODE_64)
    cs.detail = True
    selected = [f for f in funcs if 'gridfight' in f['name'].lower() or ('firefly-go' in f['name'] and any(x in f['name'].lower() for x in ('register', 'route', 'dispatch')))]
    for f in selected:
        va = int(f['va'], 16)
        raw = offset(va)
        f['calls'] = []
        asm = []
        for ins in cs.disasm(data[raw:raw+f['size']], va):
            target = None
            if ins.mnemonic in ('call', 'jmp') and ins.operands and ins.operands[0].type == X86_OP_IMM:
                target = ins.operands[0].imm
                if target in byva:
                    f['calls'].append({'at':hex(ins.address), 'target':byva[target], 'va':hex(target)})
            asm.append(f'{ins.address:#x}: {ins.mnemonic} {ins.op_str}' + (f' ; {byva[target]}' if target in byva else ''))
        f['assembly'] = asm
    output.mkdir(parents=True, exist_ok=True)
    provenance = {'sha256':hashlib.sha256(data).hexdigest(), 'bytes':len(data), 'image_base':hex(image_base), 'pclntab_offset':header, 'function_count':nfunc, 'limitations':['Compiled symbols and direct-call edges only; no source reconstruction or runtime success assertion.']}
    (output/'functions.json').write_text(json.dumps({'provenance':provenance,'functions':funcs}, indent=2), encoding='utf8')
    (output/'gridfight-workflow.json').write_text(json.dumps({'provenance':provenance,'functions':selected}, indent=2), encoding='utf8')
    print(json.dumps({**provenance,'selected_functions':len(selected)}))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--binary', type=pathlib.Path, required=True)
    parser.add_argument('--output', type=pathlib.Path, required=True)
    args = parser.parse_args()
    dump(args.binary, args.output)