#!/usr/bin/env python3
"""Recover qagame authentication announcement tables and bounded recipient paths."""
import json
import gzip
from pathlib import Path
import re
import struct

from analyze_auth_drop import ascii_literal
from analyze_deferred_auth import partition_function, trace_slice
from analyze_engine_control import control_flow
from analyze_native_dispatch import recover
from analyze_qvm import decode, write_json
from analyze_vm_extensions import native_functions, verified


def recover_table(data, header, instructions, base, lower, upper, entry, end):
    start = base + lower*4
    size = (upper-lower+1)*4
    if lower > upper or start < 0 or start+size > header['data_length']:
        raise ValueError('jump table outside initialized data')
    if instructions[entry]['opcode'] != 'ENTER' or not entry < end <= len(instructions):
        raise ValueError('invalid function interval')
    rows = []
    for value in range(lower, upper+1):
        address = base+value*4
        offset = header['data_offset']+address
        if offset+4 > len(data):
            raise ValueError('truncated jump table')
        target = struct.unpack_from('<I', data, offset)[0]
        if not entry < target < end:
            raise ValueError('jump target outside announcement body')
        rows.append({'value': value, 'vm_address': address, 'file_offset': offset,
                     'target_instruction': target, 'entry_window': instructions[target:target+12]})
    return {'base': base, 'lower': lower, 'upper': upper, 'rows': rows}


def main():
    report = json.loads(Path('docs/research/qvm/qagame.json').read_text())
    data = verified(Path(report['source']['member']['output']).read_bytes(), report['sha256'])
    header, ins = decode(data)
    function = partition_function(ins, 165352, data, header)
    end = function['linear_end_exclusive']
    tables = {'status': recover_table(data, header, ins, 29152, -1, 2, 165352, end),
              'role': recover_table(data, header, ins, 29168, -1, 8, 165352, end)}
    # Verify these constants occur in the supplied computed-jump expressions.
    for at, expected in [(165403, 29152), (165493, 29168)]:
        if ins[at]['opcode'] != 'CONST' or ins[at]['operand'] != expected:
            raise ValueError('announcement switch identity changed')
    memory = {row['vm_address']: row['target_instruction'] for table in tables.values() for row in table['rows']}
    formatting = {}
    for name, start, stop, values in [('status', 165384, 165454, range(-2, 4)),
                                     ('role', 165474, 165610, range(-2, 10))]:
        traces = []
        for value in values:
            model = trace_slice(ins, start, stop, {**memory, 0x10000+3688: value}, {2272: 0x10000})
            copies = []
            for call in model['calls']:
                if call['target'] != 69338:
                    raise ValueError('unexpected label formatter')
                address = call['args'][12]
                text = ascii_literal(data, header['data_offset']+address,
                                     header['data_offset']+header['data_length']+header['literal_length'])
                if text is None:
                    raise ValueError('missing status/role label')
                copies.append({'vm_address': address, 'text': text, 'capacity': call['args'][16]})
            traces.append({'input_value': value, 'selected_labels': copies,
                           'empty_string_writes': model['byte_writes'], 'trace': model['trace']})
        formatting[name] = traces
    recipient_traces = []
    record, target_entity, client_base = 0x10000, 0x30000, 0x50000
    for mode, target, broadcast in [(1, 0, 0), (1, target_entity, 2), (2, 0, 0),
                                    (2, target_entity, 0), (0, 0, 1), (0, 0, 2),
                                    (0, 0, 0), (3, 0, 1), (-1, 0, 2)]:
        state = {record+140: 2, record+3688: 5, 35002972: broadcast,
                 target_entity+520: client_base+3*6336, 34992496: client_base}
        model = trace_slice(ins, 165642, 165816, state, {2272: record, 2276: mode, 2280: target})
        recipient_traces.append({'mode': mode, 'third_argument': target, 'broadcast_global': broadcast,
                                 'record_number': 2, 'target_slot': 3, 'calls': model['calls'], 'trace': model['trace']})
    references, entry = [], None
    for i in ins:
        if i['opcode'] == 'ENTER':
            entry = i['index']
        if i['opcode'] == 'CONST' and i['operand'] == 3200:
            references.append({'instruction': i['index'], 'function_entry': entry,
                               'window': ins[max(entry, i['index']-5):i['index']+12]})
    output = {'source': report['source'], 'sha256': report['sha256'], 'announcement_function': function,
              'tables': tables, 'synthetic_label_traces': formatting,
              'synthetic_recipient_traces': recipient_traces, 'constant_3200_references': references,
              'support_functions': {str(e): partition_function(ins, e, data, header) for e in [36038, 5123, 2294]},
              'limitations': 'Static qagame and synthetic memory only. Calls are recorded, not entered; '
                             'their zero return values are placeholders, so formatted command pointers are not recovered outputs. '
                             'Literal displacement scanning does not recover bulk clears or alias writes; no permissions/runtime claim.'}
    output['native'] = {}
    for label in ('linux--Quake3-UrT-Ded.x86_64', 'linux--Quake3-UrT.x86_64'):
        native = json.loads(Path(f'docs/research/native/{label}.json').read_text())
        raw = verified((Path('.cache/UrbanTerror43')/native['source']['path']).read_bytes(), native['source']['sha256'])
        artifact = native['artifacts']['disassembly.txt']
        text = verified(gzip.open(artifact['path'], 'rb').read(), artifact['uncompressed_sha256']).decode()
        functions = native_functions(text)
        dispatch = recover(raw, text, 'SV_GameSystemCalls')
        case = next(c for c in dispatch['cases'] if 17 in c['numbers_sharing_target'])
        branch = next(line for line in case['first_instructions'] if '\tje\t' in line)
        target = int(re.search(r'\tje\t0x([0-9a-f]+)', branch)[1], 16)
        graph = control_flow(functions['SV_GameSystemCalls'])
        index = next(k for k, i in enumerate(graph['instructions']) if i['address'] == target)
        output['native'][label] = {'source': native['source'], 'disassembly_sha256': artifact['uncompressed_sha256'],
                                   'trap_17': case, 'minus_one_branch': graph['instructions'][index:index+6],
                                   'graphs': {name: control_flow(functions[name]) for name in ['SV_SendServerCommand', 'SV_AddServerCommand']}}
    write_json(Path('docs/research/auth-announcements.json'), output)
    print('Recovered status/role tables, eighteen label cases, nine recipient paths and flag references')


if __name__ == '__main__':
    main()