#!/usr/bin/env python3
"""Static ELF/PE/Mach-O layout, dependency, symbol and disassembly inventory."""
import argparse
from collections import Counter
import gzip
import hashlib
import json
from pathlib import Path
import platform
import re
import shutil
import struct
import subprocess
from analyze_qvm import write_json, write_rows

MAGICS = {b'\x7fELF': 'ELF', b'MZ': 'PE', b'\xcf\xfa\xed\xfe': 'Mach-O'}


def take(data, offset, length):
    if offset < 0 or length < 0 or offset+length > len(data):
        raise ValueError('out-of-file structure')
    return data[offset:offset+length]


def unpack(fmt, data, offset):
    return struct.unpack(fmt, take(data, offset, struct.calcsize(fmt)))


def cstring(data, offset):
    if not 0 <= offset < len(data):
        raise ValueError('invalid string offset')
    end = data.find(b'\0', offset)
    if end < 0:
        raise ValueError('unterminated string')
    return data[offset:end].decode('utf-8', errors='backslashreplace')


def elf(data):
    if data[4] not in (1, 2) or data[5] != 1:
        raise ValueError('only little-endian ELF32/64 supported')
    bits = 64 if data[4] == 2 else 32
    h = unpack('<HHIQQQIHHHHHH' if bits == 64 else '<HHIIIIIHHHHHH', data, 16)
    kind, machine, _, entry, phoff, shoff, _, _, phsize, phcount, shsize, shcount, names_index = h
    sfmt = '<IIQQQQIIQQ' if bits == 64 else '<IIIIIIIIII'
    if shsize != struct.calcsize(sfmt) or not 0 < names_index < shcount:
        raise ValueError('unsupported ELF section layout')
    headers = [unpack(sfmt, data, shoff+i*shsize) for i in range(shcount)]
    names = take(data, headers[names_index][4], headers[names_index][5])
    sections, symbols, dependencies, relocations = [], [], [], []
    for index, s in enumerate(headers):
        name, typ, flags, address, offset, size, link, info, align, entsize = s
        name = cstring(names, name)
        if typ != 8:
            take(data, offset, size)
        sections.append({'index': index, 'name': name, 'type': typ, 'flags': flags,
                         'address': address, 'file_offset': offset, 'bytes': size,
                         'executable': bool(flags & 4), 'zero_fill': typ == 8})
        if typ in (2, 11, 6):
            if not 0 <= link < shcount or not entsize or size % entsize:
                raise ValueError('invalid ELF linked table')
            linked = take(data, headers[link][4], headers[link][5])
            for n in range(size//entsize):
                at = offset+n*entsize
                if typ == 6:
                    tag, value = unpack('<qQ' if bits == 64 else '<iI', data, at)
                    if tag == 1:
                        dependencies.append(cstring(linked, value))
                    continue
                if bits == 64:
                    nm, flags_s, other, section, value, length = unpack('<IBBHQQ', data, at)
                else:
                    nm, value, length, flags_s, other, section = unpack('<IIIBBH', data, at)
                if nm:
                    symbols.append({'name': cstring(linked, nm), 'address': value,
                                    'bytes': length, 'section': section, 'kind': flags_s & 15,
                                    'binding': flags_s >> 4, 'visibility': other & 3,
                                    'undefined': section == 0, 'table': name})
        if typ in (4, 9):
            relocations.append({'section': name, 'entry_bytes': entsize,
                                'count': size//entsize if entsize else None})
    interpreter = None
    for i in range(phcount):
        at = phoff+i*phsize
        ph = unpack('<IIQQQQQQ' if bits == 64 else '<IIIIIIII', data, at)
        if ph[0] == 3:
            interpreter = cstring(data, ph[2] if bits == 64 else ph[1])
    return {'format': 'ELF', 'bits': bits, 'machine': machine, 'image_type': kind,
            'entry_address': entry, 'interpreter': interpreter, 'dependencies': dependencies,
            'sections': sections, 'relocations': relocations}, symbols


def pe(data):
    pe_at = unpack('<I', data, 0x3c)[0]
    if take(data, pe_at, 4) != b'PE\0\0':
        raise ValueError('MZ without valid PE signature')
    machine, count, timestamp, sym_at, sym_count, opt_len, flags = unpack('<HHIIIHH', data, pe_at+4)
    opt = take(data, pe_at+24, opt_len)
    magic = unpack('<H', opt, 0)[0]
    if magic not in (0x10b, 0x20b):
        raise ValueError('unsupported PE optional header')
    bits = 32 if magic == 0x10b else 64
    base = unpack('<I' if bits == 32 else '<Q', opt, 28 if bits == 32 else 24)[0]
    entry = unpack('<I', opt, 16)[0]
    sections = []
    for i in range(count):
        s = unpack('<8sIIIIIIHHI', data, pe_at+24+opt_len+i*40)
        name, vsize, va, size, offset, rel_at, _, rel_count, _, chars = s
        take(data, offset, size)
        sections.append({'index': i+1, 'name': name.rstrip(b'\0').decode('ascii', errors='replace'),
                         'address': base+va, 'rva': va, 'virtual_bytes': vsize,
                         'file_offset': offset, 'bytes': size, 'flags': chars,
                         'executable': bool(chars & 0x20000000), 'relocation_count': rel_count})
    def rva_offset(rva):
        for section in sections:
            delta = rva-section['rva']
            if 0 <= delta < section['bytes']:
                return section['file_offset']+delta
        raise ValueError(f'unmapped PE RVA {rva:#x}')
    directories_at = 96 if bits == 32 else 112
    directory_count = unpack('<I', opt, directories_at-4)[0]
    directories = [{'index': i, 'rva_or_security_file_offset': unpack('<II', opt, directories_at+i*8)[0],
                    'bytes': unpack('<II', opt, directories_at+i*8)[1]} for i in range(directory_count)]
    imports = []
    if directory_count > 1 and directories[1]['rva_or_security_file_offset']:
        at = rva_offset(directories[1]['rva_or_security_file_offset'])
        for _ in range(4096):
            desc = unpack('<IIIII', data, at)
            if not any(desc):
                break
            ilt, _, _, name, iat = desc
            row = {'library': cstring(data, rva_offset(name)), 'symbols': []}
            ptr = rva_offset(ilt or iat)
            for _ in range(65536):
                thunk = unpack('<I' if bits == 32 else '<Q', data, ptr)[0]
                ptr += bits//8
                if not thunk:
                    break
                if thunk & (1 << (bits-1)):
                    row['symbols'].append({'ordinal': thunk & 65535})
                else:
                    off = rva_offset(thunk)
                    row['symbols'].append({'hint': unpack('<H', data, off)[0], 'name': cstring(data, off+2)})
            else:
                raise ValueError('unterminated PE import thunks')
            imports.append(row)
            at += 20
        else:
            raise ValueError('unterminated PE import descriptors')
    symbols = []
    if sym_count:
        table = take(data, sym_at, sym_count*18)
        strings_at = sym_at+sym_count*18
        strings = take(data, strings_at, unpack('<I', data, strings_at)[0])
        i = 0
        while i < sym_count:
            name, value, section, typ, storage, aux = unpack('<8sIhHBB', table, i*18)
            nm = cstring(strings, unpack('<I', name, 4)[0]) if name[:4] == b'\0'*4 else name.rstrip(b'\0').decode('utf-8', errors='backslashreplace')
            address = sections[section-1]['address']+value if 0 < section <= count else value
            symbols.append({'name': nm, 'address': address, 'section': section, 'type': typ,
                            'storage_class': storage, 'undefined': section == 0,
                            'coff_index': i, 'auxiliary_entries': aux})
            i += 1+aux
        if i != sym_count:
            raise ValueError('invalid COFF auxiliary count')
    return {'format': 'PE', 'bits': bits, 'machine': machine, 'image_base': base,
            'entry_address': base+entry, 'entry_rva': entry, 'coff_timestamp_untrusted': timestamp,
            'characteristics': flags, 'data_directories': directories,
            'dependencies': [row['library'] for row in imports], 'imports': imports,
            'sections': sections, 'coff_entry_count_including_aux': sym_count,
            'section_tail_bytes': len(data)-max(s['file_offset']+s['bytes'] for s in sections)}, symbols


def macho(data):
    _, cpu, subtype, kind, ncmds, cmd_bytes, flags, _ = unpack('<8I', data, 0)
    commands_end = 32+cmd_bytes
    take(data, 32, cmd_bytes)
    at = 32
    sections, symbols, dependencies, commands, identity, rpaths = [], [], [], [], None, []
    symtab = None
    for _ in range(ncmds):
        cmd, size = unpack('<II', data, at)
        if size < 8 or at+size > commands_end:
            raise ValueError('invalid Mach-O load command')
        commands.append({'command': hex(cmd), 'file_offset': at, 'bytes': size})
        command = take(data, at, size)
        if cmd == 0x19:
            _, _, seg, va, vsize, off, length, maxprot, initprot, count, _ = unpack('<II16sQQQQIIII', command, 0)
            take(data, off, length)
            for i in range(count):
                s = unpack('<16s16sQQIIIIIIII', command, 72+i*80)
                name, segment, address, length, offset, _, reloc_at, reloc_count, sf, _, _, _ = s
                zero = sf & 255 in (1, 12, 18)
                if not zero:
                    take(data, offset, length)
                sections.append({'index': len(sections)+1, 'name': name.rstrip(b'\0').decode(),
                                 'segment': segment.rstrip(b'\0').decode(), 'address': address,
                                 'file_offset': offset, 'bytes': length, 'zero_fill': zero, 'flags': sf,
                                 'executable': bool(sf & 0x80000400), 'relocation_count': reloc_count})
        elif cmd in (12, 13, 0x80000018, 0x8000001f, 0x20, 0x80000023):
            name_at, _, current, compatibility = unpack('<4I', command, 8)
            row = {'name': cstring(command, name_at), 'current_version': current,
                   'compatibility_version': compatibility, 'command': hex(cmd)}
            if cmd == 13:
                identity = row
            else:
                dependencies.append(row)
        elif cmd == 0x8000001c:
            rpaths.append(cstring(command, unpack('<I', command, 8)[0]))
        elif cmd == 2:
            symtab = unpack('<4I', command, 8)
        at += size
    if at != commands_end:
        raise ValueError('Mach-O command size mismatch')
    if symtab:
        offset, count, str_at, str_bytes = symtab
        strings = take(data, str_at, str_bytes)
        for i in range(count):
            name, typ, section, desc, address = unpack('<IBBHQ', data, offset+i*16)
            if name:
                symbols.append({'name': cstring(strings, name), 'address': address, 'section': section,
                                'type': typ, 'description': desc, 'external': bool(typ & 1),
                                'debug_stab': bool(typ & 0xe0), 'undefined': (typ & 0xe) == 0})
    return {'format': 'Mach-O', 'bits': 64, 'machine': cpu, 'subtype': subtype,
            'image_type': kind, 'flags': flags, 'load_commands': commands,
            'sections': sections, 'dependencies': dependencies, 'library_identity': identity,
            'rpaths': rpaths}, symbols


def parse(data):
    if data.startswith(b'\x7fELF'):
        return elf(data)
    if data.startswith(b'MZ'):
        return pe(data)
    if data.startswith(b'\xcf\xfa\xed\xfe'):
        return macho(data)
    raise ValueError('unsupported native format (fat/big-endian images need separate support)')


def ascii_strings(data):
    return [{'file_offset': m.start(), 'encoding': 'ASCII', 'text': m.group().decode('ascii')}
            for m in re.finditer(rb'[\x20-\x7e]{4,}', data)]


def run_tool(args):
    result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=120)
    if result.returncode:
        raise ValueError(f'{args[0]} failed: {result.stderr.decode(errors="replace")[:500]}')
    return result.stdout


def write_gzip(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('wb') as raw, gzip.GzipFile(filename='', fileobj=raw, mode='wb', mtime=0) as out:
        out.write(data)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', type=Path, default=Path('.cache/UrbanTerror43'))
    parser.add_argument('--reports', type=Path, default=Path('docs/research/native'))
    parser.add_argument('--artifacts', type=Path, default=Path('assets/analysis/native'))
    parser.add_argument('--objdump', default='llvm-objdump' if shutil.which('llvm-objdump') else 'objdump')
    args = parser.parse_args()
    snapshot = json.loads(Path('docs/inventory/snapshot.json').read_text())
    actual = {p.relative_to(args.source).as_posix() for p in args.source.rglob('*') if p.is_file()}
    if actual != {f['path'] for f in snapshot['files']}:
        raise ValueError('installation file set changed')
    targets, scan = [], []
    member_count = 0
    for record in snapshot['files']:
        path = args.source/record['path']
        if path.is_symlink():
            raise ValueError('source symlink forbidden')
        data = path.read_bytes()
        if len(data) != record['bytes'] or hashlib.sha256(data).hexdigest() != record['sha256']:
            raise ValueError(f'{path}: snapshot mismatch')
        fmt = next((name for magic, name in MAGICS.items() if data.startswith(magic)), None)
        scan.append({'path': record['path'], 'format': fmt, 'sha256': record['sha256'], 'origin': 'installation'})
        if fmt:
            targets.append((path, record, data))
    # Exhaustively inspect every extracted member prefix, independent of suffix.
    for manifest in sorted(Path('docs/extraction/archives').glob('*.jsonl')):
        for line in manifest.read_text().splitlines()[1:]:
            row = json.loads(line)
            if not row['output']:
                continue
            member_count += 1
            path = Path(row['output'])
            if path.is_symlink():
                raise ValueError('extracted symlink forbidden')
            with path.open('rb') as source:
                prefix = source.read(4)
            fmt = next((name for magic, name in MAGICS.items() if prefix.startswith(magic)), None)
            if fmt or path.suffix.lower() in {'.dll', '.exe', '.so', '.dylib'}:
                scan.append({'path': row['output'], 'format': fmt, 'sha256': row['sha256'], 'origin': 'archive'})
                raise ValueError('new native archive member discovered: extend target registry before proceeding')
    if len(targets) != 8:
        raise ValueError('expected eight supplied native images; review discovery changes')
    objdump = shutil.which(args.objdump)
    if not objdump:
        raise ValueError('LLVM objdump required for reproducible disassembly')
    tool_version = run_tool([objdump, '--version']).decode().splitlines()[:2]
    summaries = []
    library_usage = []
    anchors = re.compile(r'^(?:_)?(?:VM_|CL_CgameSystemCalls|SV_GameSystemCalls|CL_UISystemCalls|Com_Frame|SV_Frame|CL_Frame|FS_Startup|FS_FOpenFileRead|Sys_LoadDll|SDL_GetVersion|SDL_Linked_Version|SDL_Init|SDL_Quit)')
    for path, record, data in targets:
        layout, symbols = parse(data)
        label = {'ELF': 'linux', 'PE': 'windows', 'Mach-O': 'macos'}[layout['format']]+'--'+path.name
        sdl_symbols = [s for s in symbols if not s.get('debug_stab') and s['name'].lstrip('_').startswith('SDL_')]
        library_usage.append({'image': label, 'sha256': record['sha256'],
                              'imported_SDL_names': sorted({s['name'].lstrip('_') for s in sdl_symbols if s['undefined']}),
                              'defined_SDL_names': sorted({s['name'].lstrip('_') for s in sdl_symbols if not s['undefined']}),
                              'external_defined_SDL_names': sorted({s['name'].lstrip('_') for s in sdl_symbols if not s['undefined'] and s.get('external', s.get('binding') == 1 or s.get('storage_class') == 2)})})
        strings = ascii_strings(data)
        strings += [{'file_offset': m.start(), 'encoding': 'UTF-16LE printable candidate',
                     'text': m.group().decode('utf-16le')}
                    for m in re.finditer(rb'(?:[\x20-\x7e]\x00){4,}', data)]
        for row in strings:
            section = next((s for s in layout['sections'] if not s.get('zero_fill') and s['file_offset'] <= row['file_offset'] < s['file_offset']+s['bytes']), None)
            row['section'] = section['name'] if section else None
            row['address'] = section['address']+row['file_offset']-section['file_offset'] if section else None
        library_strings = [s for s in strings if re.search(r'\.dll\b|\.dylib\b|\.so(?:\.[0-9]+)*\b', s['text'], re.I)]
        build_strings = [s for s in strings if re.search(r'ioQ3|GCC:|clang version|SDL 1\.|urt 4\.3|GLIBC_|MinGW', s['text'])]
        embedded = [{'file_offset': m.start(), 'signature': magic.hex(), 'status': 'unvalidated signature candidate'}
                    for magic in MAGICS for m in re.finditer(re.escape(magic), data) if m.start()]
        for candidate in embedded:
            try:
                nested, _ = parse(data[candidate['file_offset']:])
                candidate.update(status='structurally parseable candidate; embedded role unproven', format=nested['format'])
            except (ValueError, IndexError, struct.error) as exc:
                candidate.update(status='rejected by supported native structural parser', reason=str(exc))
        summary = {'schema_version': 1, 'source': record, 'layout': layout,
                   'symbol_entries': len(symbols), 'undefined_symbols': sum(s['undefined'] for s in symbols),
                   'defined_function_symbols': sum(not s['undefined'] and not s.get('debug_stab', False) and ((s.get('kind') == 2) if layout['format'] == 'ELF' else (s.get('type', 0) == 32) if layout['format'] == 'PE' else s.get('section') in {sec['index'] for sec in layout['sections'] if sec['executable']}) for s in symbols),
                   'debug_stab_symbol_entries': sum(s.get('debug_stab', False) for s in symbols),
                   'debug_section_names': [s['name'] for s in layout['sections'] if s['name'].startswith(('.debug', '__debug'))],
                   'anchors': [s for s in symbols if anchors.match(s['name']) and not s['undefined'] and not s.get('debug_stab', False)],
                   'library_string_candidates': library_strings, 'build_string_evidence': build_strings,
                   'embedded_header_candidates': embedded,
                   'ascii_string_count': sum(s['encoding'] == 'ASCII' for s in strings),
                   'utf16_string_candidates': sum(s['encoding'] != 'ASCII' for s in strings)}
        write_json(args.reports/(label+'.json'), summary)
        write_rows(args.reports/(label+'.symbols.jsonl'), symbols)
        write_rows(args.artifacts/(label+'.strings.jsonl.gz'), strings)
        headers = run_tool([objdump, '--macho', '--private-headers', '--section-headers', '--syms', str(path)] if layout['format'] == 'Mach-O' else [objdump, '--all-headers', str(path)])
        disassembly = run_tool([objdump, '--disassemble', str(path)])
        if path.suffix == '.dylib':
            version = next((s for s in symbols if s['name'] == '_version' and not s['debug_stab']), None)
            if version:
                section = next(s for s in layout['sections'] if not s['zero_fill'] and s['address'] <= version['address'] < s['address']+s['bytes'])
                off = section['file_offset']+version['address']-section['address']
                summary['SDL_version_data'] = {'address': version['address'], 'file_offset': off,
                                                'bytes': list(take(data, off, 3)),
                                                'status': 'static version object; linked-version accessor must be cross-checked'}
        write_gzip(args.artifacts/(label+'.headers.txt.gz'), headers)
        write_gzip(args.artifacts/(label+'.disassembly.txt.gz'), disassembly)
        summary['artifacts'] = {name: {'path': str(args.artifacts/(label+'.'+name+'.gz')), 'uncompressed_sha256': hashlib.sha256(payload).hexdigest(), 'uncompressed_bytes': len(payload)} for name, payload in [('headers.txt', headers), ('disassembly.txt', disassembly)]}
        write_json(args.reports/(label+'.json'), summary)
        summaries.append({'path': record['path'], 'sha256': record['sha256'], 'format': layout['format'],
                          'bits': layout['bits'], 'symbols': len(symbols), 'functions': summary['defined_function_symbols'],
                          'imports': layout['dependencies'], 'debug_sections': summary['debug_section_names'],
                          'debug_stab_entries': summary['debug_stab_symbol_entries']})
    write_json(args.reports/'summary.json', {'images': summaries, 'objdump_version': tool_version,
               'environment': {'os': platform.system(), 'architecture': platform.machine(), 'python': platform.python_version()},
               'installation_files_verified': len(scan), 'extracted_member_prefixes_scanned': member_count,
               'native_archive_members': 0, 'scope': 'All installation files and all extracted member prefixes; nonzero embedded magic hits are unvalidated candidates, not recovered libraries.'})
    write_json(args.reports/'installation-scan.json', scan)
    write_json(args.reports/'library-usage.json', {'images': library_usage,
               'limitations': 'Static import/definition names; not proof of runtime reachability or behavior. Mach-O external symbols are export candidates, not a full dyld bind/export analysis.'})
    print('Analyzed eight native images; verified all 39 originals and scanned all extracted member prefixes')


if __name__ == '__main__':
    main()