#!/usr/bin/env python3
"""Bounded Mach-O STABS and PE resource recovery; no native code execution."""
from collections import Counter
import hashlib
import json
from pathlib import Path
import struct
from analyze_native import parse, take, unpack, cstring
from analyze_qvm import write_json, write_rows
from analyze_vm_extensions import verified

STABS = {0x20: 'N_GSYM', 0x22: 'N_FNAME', 0x24: 'N_FUN', 0x26: 'N_STSYM',
         0x28: 'N_LCSYM', 0x2e: 'N_BNSYM', 0x32: 'N_AST', 0x3c: 'N_OPT',
         0x40: 'N_RSYM', 0x44: 'N_SLINE', 0x4e: 'N_ENSYM', 0x60: 'N_SSYM',
         0x64: 'N_SO', 0x66: 'N_OSO', 0x80: 'N_LSYM', 0x82: 'N_BINCL',
         0x84: 'N_SOL', 0x86: 'N_PARAMS', 0x88: 'N_VERSION', 0x8a: 'N_OLEVEL',
         0xa0: 'N_PSYM', 0xa2: 'N_EINCL', 0xa4: 'N_ENTRY', 0xc0: 'N_LBRAC',
         0xc2: 'N_EXCL', 0xe0: 'N_RBRAC', 0xe2: 'N_BCOMM', 0xe4: 'N_ECOMM'}


def stab_records(data, layout):
    if layout['format'] != 'Mach-O' or layout['bits'] != 64:
        raise ValueError('STABS reader supports thin Mach-O64 only')
    cmd = next(c for c in layout['load_commands'] if c['command'] == '0x2')
    offset, count, strings_at, strings_bytes = unpack('<4I', data, cmd['file_offset']+8)
    strings = take(data, strings_at, strings_bytes)
    take(data, offset, count*16)
    rows = []
    for i in range(count):
        name, typ, section, desc, value = unpack('<IBBHQ', data, offset+i*16)
        if typ & 0xe0:
            rows.append({'symbol_index': i, 'file_offset': offset+i*16, 'name': cstring(strings, name),
                         'type': typ, 'kind': STABS.get(typ, 'UNKNOWN'), 'section': section,
                         'description': desc, 'value': value})
    return rows


def decode_stabs(data, layout, rows):
    sources, objects, functions, unresolved = [], [], [], []
    directory = source = obj = None
    active = None
    for row in rows:
        kind, name = row['kind'], row['name']
        if kind == 'N_SO':
            if name.endswith('/'):
                directory = name
            elif name:
                source = name
                sources.append({'directory': directory, 'source': source, 'record': row['symbol_index']})
            else:
                source = obj = None
                directory = None
        elif kind == 'N_OSO':
            obj = name
            objects.append({'object': obj, 'mtime_untrusted': row['value'], 'record': row['symbol_index']})
        elif kind == 'N_FUN' and name:
            if active:
                unresolved.append({'reason': 'function without closing N_FUN', 'record': active['record']})
            active = {'name': name, 'address': row['value'], 'section': row['section'],
                      'directory': directory, 'source': source, 'object': obj, 'record': row['symbol_index']}
        elif kind == 'N_FUN' and not name:
            if not active:
                unresolved.append({'reason': 'closing N_FUN without open function', 'record': row['symbol_index']})
                continue
            size = row['value']
            section = next((s for s in layout['sections'] if s['index'] == active['section'] and
                            s['executable'] and not s['zero_fill'] and s['address'] <= active['address'] and
                            active['address']+size <= s['address']+s['bytes']), None)
            if section is None or size == 0:
                unresolved.append({'reason': 'function range is not nonempty file-backed executable code',
                                   'record': active['record'], 'size': size})
            else:
                offset = section['file_offset']+active['address']-section['address']
                functions.append({**active, 'bytes': size, 'file_offset': offset,
                                  'sha256': hashlib.sha256(take(data, offset, size)).hexdigest(),
                                  'end_record': row['symbol_index']})
            active = None
        elif kind == 'UNKNOWN':
            unresolved.append({'reason': 'unknown STABS type', 'record': row['symbol_index']})
    if active:
        unresolved.append({'reason': 'unclosed final function', 'record': active['record']})
    return {'record_count': len(rows), 'kind_counts': dict(sorted(Counter(r['kind'] for r in rows).items())),
            'sources': sources, 'objects': objects, 'functions': functions, 'unresolved': unresolved,
            'limitations': 'N_FUN end values are interpreted as byte lengths for this linker debug-map pattern. '
            'Source/object paths and timestamps are embedded build evidence, not verified source identities. '
            'Raw records retain begin/end/include/static symbols; source types are not reconstructed from absent records.'}


def rva_range(data, layout, rva, size):
    matches = [s for s in layout['sections'] if s['rva'] <= rva and rva+size <= s['rva']+s['bytes']]
    if len(matches) != 1:
        raise ValueError('resource range is not uniquely file-backed')
    offset = matches[0]['file_offset']+rva-matches[0]['rva']
    return offset, take(data, offset, size)


def resource_leaves(data, layout, regions=None):
    if layout['format'] != 'PE':
        raise ValueError('resource reader expects PE')
    directory = next((d for d in layout['data_directories'] if d['index'] == 2), None)
    if not directory or not directory['bytes']:
        return []
    base = directory['rva_or_security_file_offset']
    tree_file_offset, tree = rva_range(data, layout, base, directory['bytes'])
    def record(offset, size, kind):
        if regions is not None:
            regions.append({'file_offset': tree_file_offset+offset, 'bytes': size, 'kind': kind})
    leaves = []
    nodes = 0
    def walk(offset, path, ancestors):
        nonlocal nodes
        nodes += 1
        if offset in ancestors or len(path) > 8 or nodes > 65536:
            raise ValueError('cyclic or excessive resource tree')
        header = unpack('<IIHHHH', tree, offset)
        named, ids = header[-2:]
        entries = take(tree, offset+16, (named+ids)*8)
        record(offset, 16+(named+ids)*8, 'resource_directory_and_entries')
        for i in range(named+ids):
            name, target = unpack('<II', entries, i*8)
            if bool(name & 0x80000000) != (i < named):
                raise ValueError('resource named/ID entry ordering mismatch')
            if name & 0x80000000:
                at = name & 0x7fffffff
                length = unpack('<H', tree, at)[0]
                component = take(tree, at+2, length*2).decode('utf-16le')
                record(at, 2+length*2, 'resource_name_utf16')
            else:
                component = name
            current = path+[component]
            if target & 0x80000000:
                walk(target & 0x7fffffff, current, ancestors | {offset})
            else:
                rva, size, codepage, reserved = unpack('<4I', tree, target)
                record(target, 16, 'resource_data_descriptor')
                if reserved:
                    raise ValueError('nonzero resource data reserved field')
                file_offset, payload = rva_range(data, layout, rva, size)
                if regions is not None:
                    regions.append({'file_offset': file_offset, 'bytes': size, 'kind': 'resource_payload'})
                leaves.append({'path': current, 'data_entry_relative_offset': target,
                               'rva': rva, 'file_offset': file_offset, 'bytes': size, 'codepage': codepage,
                               'sha256': hashlib.sha256(payload).hexdigest()})
    walk(0, [], set())
    return leaves


def icon_group(payload):
    reserved, kind, count = unpack('<3H', payload, 0)
    if reserved or kind != 1 or len(payload) != 6+14*count:
        raise ValueError('invalid group icon header/length')
    rows = []
    for i in range(count):
        width, height, colors, reserved, planes, bits, size, ident = unpack('<4B2HIH', payload, 6+i*14)
        if reserved:
            raise ValueError('invalid icon entry reserved field')
        rows.append({'width': width or 256, 'height': height or 256, 'colors': colors,
                     'planes': planes, 'bits': bits, 'bytes': size, 'resource_id': ident})
    return rows


def string_table(payload, block):
    if not isinstance(block, int) or block < 1:
        raise ValueError('string table requires a positive numeric block ID')
    rows, offset = [], 0
    for i in range(16):
        length = unpack('<H', payload, offset)[0]
        offset += 2
        value = take(payload, offset, length*2).decode('utf-16le')
        offset += length*2
        rows.append({'id': (block-1)*16+i, 'text': value})
    if offset != len(payload):
        raise ValueError('trailing string table bytes')
    return rows


def icon_dib(payload):
    # Retain surplus bytes rather than assuming Windows ignores them.
    size, width, doubled_height, planes, bits, compression, image_bytes, _, _, colors, _ = unpack('<IiiHHIIiiII', payload, 0)
    if size != 40 or width <= 0 or doubled_height <= 0 or doubled_height % 2 or planes != 1 or bits not in (8, 32) or compression:
        raise ValueError('unsupported icon DIB layout')
    height = doubled_height//2
    palette = (colors or (1 << bits if bits <= 8 else 0))*4
    xor = ((width*bits+31)//32)*4*height
    mask = ((width+31)//32)*4*height
    expected = 40+palette+xor+mask
    take(payload, 0, expected)
    return {'width': width, 'height': height, 'bits': bits, 'declared_image_bytes': image_bytes,
            'palette_bytes': palette, 'xor_bytes': xor, 'standard_and_mask_bytes': mask,
            'standard_layout_bytes': expected, 'surplus_bytes': len(payload)-expected,
            'surplus_sha256': hashlib.sha256(payload[expected:]).hexdigest() if len(payload)>expected else None}


def reconstruct_icon(data, leaves, group):
    rows = icon_group(take(data, group['file_offset'], group['bytes']))
    header = bytearray(struct.pack('<3H', 0, 1, len(rows)))
    images = []
    offset = 6+16*len(rows)
    for row in rows:
        matches = [r for r in leaves if r['path'] == [3, row['resource_id'], group['path'][2]]]
        if len(matches) != 1 or matches[0]['bytes'] != row['bytes']:
            raise ValueError('missing/ambiguous icon language or mismatched image size')
        leaf = matches[0]
        image = take(data, leaf['file_offset'], leaf['bytes'])
        header.extend(struct.pack('<4B2H2I', row['width'] % 256, row['height'] % 256,
                                  row['colors'], 0, row['planes'], row['bits'], len(image), offset))
        images.append(image)
        offset += len(image)
    return bytes(header)+b''.join(images), rows


def main():
    summaries = []
    for label in ('macos--libSDL-1.2.0.dylib', 'macos--Quake3-UrT.x86_64',
                  'windows--Quake3-UrT.exe', 'windows--Quake3-UrT-Ded.exe'):
        report = json.loads(Path(f'docs/research/native/{label}.json').read_text())
        data = verified((Path('.cache/UrbanTerror43')/report['source']['path']).read_bytes(), report['source']['sha256'])
        layout, _ = parse(data)
        output = {'source': report['source']}
        if layout['format'] == 'Mach-O':
            rows = stab_records(data, layout)
            output.update(decode_stabs(data, layout, rows))
            artifact = Path(f'assets/analysis/native/{label}.stabs.jsonl.gz')
            write_rows(artifact, rows)
            output['raw_records'] = str(artifact)
            summaries.append({'image': label, 'stabs': len(rows), 'function_ranges': len(output['functions']),
                              'objects': len(output['objects']), 'unresolved': len(output['unresolved'])})
        else:
            leaves = resource_leaves(data, layout)
            root = Path(f'assets/analysis/native/resources/{label}')
            icons = []
            for i, leaf in enumerate(leaves):
                path = root/f'{i:06d}.bin'
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_bytes(take(data, leaf['file_offset'], leaf['bytes']))
                leaf['output'] = str(path)
                if len(leaf['path']) == 3 and leaf['path'][0] == 3:
                    leaf['dib'] = icon_dib(take(data, leaf['file_offset'], leaf['bytes']))
                if len(leaf['path']) == 3 and leaf['path'][0] == 6:
                    leaf['strings'] = string_table(take(data, leaf['file_offset'], leaf['bytes']), leaf['path'][1])
                if len(leaf['path']) == 3 and leaf['path'][0] == 14:
                    ico, entries = reconstruct_icon(data, leaves, leaf)
                    path = root/f'{i:06d}.ico'
                    path.write_bytes(ico)
                    icons.append({'group_path': leaf['path'], 'entries': entries, 'output': str(path),
                                  'bytes': len(ico), 'sha256': hashlib.sha256(ico).hexdigest()})
            output.update({'leaves': leaves, 'icon_groups': icons, 'limitations':
                           'All leaves are preserved. Shipped icon groups, DIB headers/extents and string tables are decoded. '
                           'Pixel/mask behavior and surplus icon bytes remain unclassified; other resource types are not decoded.'})
            summaries.append({'image': label, 'resource_leaves': len(leaves),
                              'type_counts': dict(Counter(str(r['path'][0]) for r in leaves)), 'icon_groups': len(icons)})
        write_json(Path(f'docs/research/native/{label}.metadata.json'), output)
    write_json(Path('docs/research/native/metadata-summary.json'), summaries)
    print(json.dumps(summaries, indent=2))


if __name__ == '__main__':
    main()