#!/usr/bin/env python3
"""Preserve the auth-drop call boundary and its native invalid-index continuation."""
import gzip
import json
from pathlib import Path
import re
from analyze_native import parse, take
from analyze_native_dispatch import recover
from analyze_qvm import decode, write_json
from analyze_vm_extensions import verified, native_functions


def ascii_literal(data, offset, end):
    if not 0 <= offset < end <= len(data):
        return None
    stop = data.find(b'\0', offset, end)
    if stop < 0:
        return None
    payload = data[offset:stop]
    if len(payload) < 4 or not all(32 <= c <= 126 or c in (9, 10, 13) for c in payload):
        return None
    return payload.decode('ascii')


def invalid_continuation(case, lines):
    # Restrict this recognizer to the actual signed-index guard pattern.
    guards = [re.search(r'\t(js|jge)\s+0x([0-9a-f]+)', line) for line in case['first_instructions']]
    guards = [g for g in guards if g]
    if len(guards) != 2 or [g[1] for g in guards] != ['js', 'jge'] or guards[0][2] != guards[1][2]:
        raise ValueError('unrecognized auth-drop index guards')
    address = int(guards[0][2], 16)
    indexes = {}
    for i, line in enumerate(lines):
        m = re.match(r'\s*([0-9a-f]+):', line)
        if m:
            indexes[int(m[1], 16)] = i
    if address not in indexes:
        raise ValueError('guard continuation is not disassembled')
    index = indexes[address]
    continuation = lines[index:index+2]
    if len(continuation) != 2 or not re.search(r'\txorl\s+%eax, %eax', continuation[0]) or not re.search(r'\tjmp\s+0x', continuation[1]):
        raise ValueError('invalid index does not match known zero-result continuation')
    target = int(re.search(r'\tjmp\s+0x([0-9a-f]+)', continuation[1])[1], 16)
    if target not in indexes:
        raise ValueError('missing common return continuation')
    start = indexes[target]
    epilogue = []
    for line in lines[start:start+20]:
        epilogue.append(line.strip())
        if re.search(r'\tretq?\s*$', line):
            break
        m = re.search(r'\t(\w+)\s*(.*)', line)
        if not m or m[1] not in ('movq', 'addq', 'popq', 'nop', 'nopl', 'nopw') or '%eax' in m[2] or '%rax' in m[2]:
            raise ValueError('common return could alter result or control flow')
    else:
        raise ValueError('common return not found within bound')
    return {'address': address, 'instructions': [l.strip() for l in continuation],
            'common_return_address': target, 'common_return_instructions': epilogue,
            'result': 0}


def main():
    output = {'native': {}, 'limitations': 'Linux ELF64 static boundary evidence. Literal operands are candidates, '
              'not whole-program dataflow. No runtime reference fixtures or all-platform contract.'}
    for label in ('linux--Quake3-UrT-Ded.x86_64', 'linux--Quake3-UrT.x86_64'):
        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)
        artifact = report['artifacts']['disassembly.txt']
        text = verified(gzip.open(artifact['path'], 'rb').read(), artifact['uncompressed_sha256']).decode()
        functions = native_functions(text)
        dispatch = recover(data, text, 'SV_GameSystemCalls')
        case = next(c for c in dispatch['cases'] if c['numbers_sharing_target'] == [603])
        literals = []
        function = functions['SV_Auth_DropClient']
        for line in function['lines']:
            for match in re.finditer(r'\$0x([0-9a-f]+)', line):
                address = int(match[1], 16)
                section = next((s for s in layout['sections'] if s['flags'] & 2 and not s['zero_fill'] and not s['executable'] and
                                s['address'] <= address < s['address']+s['bytes']), None)
                if section:
                    offset = section['file_offset']+address-section['address']
                    value = ascii_literal(data, offset, section['file_offset']+section['bytes'])
                    if value:
                        literals.append({'address': address, 'file_offset': offset, 'text': value,
                                         'instruction': line.strip()})
        output['native'][label] = {'source': report['source'], 'disassembly_sha256': artifact['uncompressed_sha256'],
                                   'case': case, 'invalid_index': invalid_continuation(case, text.splitlines()),
                                   'function': function, 'literal_operand_candidates': literals}
    report = json.loads(Path('docs/research/qvm/qagame.json').read_text())
    data = verified(Path(report['source']['member']['output']).read_bytes(), report['sha256'])
    header, instructions = decode(data)
    site = next(i for i in range(1, len(instructions)) if instructions[i]['opcode'] == 'CALL' and
                instructions[i-1]['opcode'] == 'CONST' and instructions[i-1]['operand'] == -604)
    start = next(i for i in range(site, -1, -1) if instructions[i]['opcode'] == 'ENTER')
    end = next(i for i in range(site+1, len(instructions)) if instructions[i]['opcode'] == 'ENTER')
    strings = []
    for ins in instructions[start:end]:
        if ins['opcode'] == 'CONST' and instructions[ins['index']+1]['opcode'] == 'ARG' and 0 <= ins['operand'] < header['data_length']+header['literal_length']:
            offset = header['data_offset']+ins['operand']
            value = ascii_literal(data, offset, header['data_offset']+header['data_length']+header['literal_length'])
            if value:
                strings.append({'instruction': ins['index'], 'vm_address': ins['operand'], 'text': value})
    output['qvm'] = {'source': report['source'], 'sha256': report['sha256'], 'call_instruction': site,
                     'containing_entry': start, 'instructions': instructions[start:end],
                     'literal_operand_candidates': strings}
    write_json(Path('docs/research/auth-drop.json'), output)
    print('Recovered both auth-drop invalid-index zero-return paths and literal operands')


if __name__ == '__main__':
    main()