#!/usr/bin/env python3
"""Scalar wrapper boundaries and native OOB data overflow compositions."""
import copy
import gzip
import itertools
import json
from pathlib import Path
from analyze_area_queries import site
from analyze_bsp_references import global_address, trace_fields
from analyze_collision_composition import get_word, put_word
from analyze_engine_control import control_flow
from analyze_message_byte_codec import MSG, DATA
from analyze_message_widths import reader, run, writer
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions, verified

SOURCE, WRAPPER_STACK=0xd0000,0xc0000


def read_prefix(g):
    g=copy.deepcopy(g)
    for x in g['instructions']:
        if x['opcode'] in ['pushq','popq']:x['opcode']='nop';x['operands']=''
    rows=g['instructions'];i=next(i for i,x in enumerate(rows) if '<MSG_ReadBits>' in x['operands'])
    t=run(g,rows[0]['address'],[rows[i]['address']],{}, {'%rdi':MSG})
    return g,i,t


def read_tail(g,i,registers,m,value):
    regs=dict(registers);regs['rax']=value
    t=run(g,g['instructions'][i+1]['address'],[x['address'] for x in g['instructions'] if x['opcode']=='retq'],m,regs)
    return t['registers']['rax'],t['trace']


def message_memory(capacity,flag=0):
    m={MSG+i:0 for i in range(40)}
    for i in range(64):m[DATA+i]=0xa5
    m[DATA-1]=m[DATA+64]=0x7b
    put_word(m,MSG+4,flag);put_word(m,MSG+8,1);put_word(m,MSG+16,DATA,8);put_word(m,MSG+24,capacity);put_word(m,MSG+36,7)
    return m


def main():
    out={'native':{},'limitations':'Synthetic Linux ELF64 instruction compositions. Original processes and native callees never execute. Scalar reader tails use supplied ReadBits return/count; separate ample/overflow scalar OOB compositions and data OOB compositions evaluate underlying writer/reader instructions. PUSH/POP and frame saves/restores omitted/stopped, caller registers explicitly retained. Real packets/compressed data loops/EOF access safety/legal cursor/allowoverflow/runtime/all-platform conformance remain open.'}
    widths={'Byte':8,'Char':8,'Short':16,'Long':32}
    for label in ['linux--Quake3-UrT-Ded.x86_64','linux--Quake3-UrT.x86_64']:
        r=json.loads(Path(f'docs/research/native/{label}.json').read_text());verified((Path('.cache/UrbanTerror43')/r['source']['path']).read_bytes(),r['source']['sha256']);a=r['artifacts']['disassembly.txt'];fs=native_functions(verified(gzip.open(a['path'],'rb').read(),a['uncompressed_sha256']).decode());names=['MSG_Write'+s for s in widths]+['MSG_Read'+s for s in widths]+['MSG_WriteData','MSG_WriteBits','MSG_ReadBits'];graphs={n:control_flow(fs[n]) for n in names}
        g=graphs['MSG_WriteBits'];old=global_address(next(x for x in g['instructions'] if x['opcode']=='addl' and '(%rip)' in x['operands']));over=global_address(next(x for x in g['instructions'] if x['opcode']=='addl' and '$0x1,' in x['operands'] and '(%rip)' in x['operands']))
        scalar=[];tails=[]
        for suffix,width in widths.items():
            wg=graphs['MSG_Write'+suffix];jump=next(x for x in wg['instructions'] if x['opcode']=='jmp');rg,ci,prefix=read_prefix(graphs['MSG_Read'+suffix]);args=[prefix['registers'][k] for k in ['rdi','rsi']]
            for value,capacity in itertools.product([0,127,0x8196a5ff,0xffffffff],[3,64]):
                forward=run(wg,wg['instructions'][0]['address'],[jump['address']],{}, {'%rdi':MSG,'%rsi':value});wargs=[forward['registers'][k] for k in ['rdi','rsi','rdx']]
                m=message_memory(capacity);put_word(m,old,0);put_word(m,over,0)
                m,requests,error,wt=writer(graphs,m,wargs[2],wargs[1]);size=get_word(m,MSG+28);bit=get_word(m,MSG+36);flag=get_word(m,MSG+4);counter=get_word(m,over);account=get_word(m,old)
                payload=[m[DATA+i] for i in range(size)];put_word(m,MSG+32,0);put_word(m,MSG+36,7)
                # Padded backing bytes allow observing the post-read EOF guard,
                # not claiming that an out-of-message read is memory safe.
                m,raw,read_error,rt=reader(graphs,m,args[1]);result,tail=read_tail(rg,ci,prefix['registers'],m,raw)
                scalar.append({'suffix':suffix,'value_word':value,'capacity':capacity,'write_arguments':wargs,'read_arguments':args,'write_error':error,'read_error':read_error,'write_requests':requests,'payload_bytes':payload,'write_size':size,'write_bit':bit,'read_count':get_word(m,MSG+32),'read_bit':get_word(m,MSG+36),'raw_read_word':raw,'result_word':result,'overflow_flag':flag,'range_counter':counter,'oldsize_word':account,'canaries':[m[DATA-1],m[DATA+64]],**trace_fields({'trace':forward['trace']+wt+prefix['trace']+rt+tail},False)})
            for size,count,value in itertools.product([0,1,32,0xffffffff],[0,1,33,0x80000000,0xffffffff],[0,127,128,255,32767,32768,0x8196a5ff,0xffffffff]):
                m={};put_word(m,MSG+28,size);put_word(m,MSG+32,count);result,trace=read_tail(rg,ci,prefix['registers'],m,value)
                tails.append({'suffix':suffix,'size_word':size,'post_count_word':count,'supplied_read_word':value,'result_word':result,'read_arguments':args,**trace_fields({'trace':trace},False)})
        data=[];dg=graphs['MSG_WriteData'];rows=dg['instructions'];start=site(dg,'testl','%edx, %edx');end=site(dg,'addq','$0x8, %rsp');calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'}
        for length,capacity,flag in itertools.product([-1,0,1,2,3,4,5,6,7,8,9,16,17,31,32],[0,3,4,8,40],[0,1]):
            m=message_memory(capacity,flag);source=[(129+37*i)&255 for i in range(32)]
            for i,v in enumerate(source):m[SOURCE+i]=v
            m[SOURCE-1]=m[SOURCE+32]=0x7b;put_word(m,old,5);put_word(m,over,0)
            regs={'%r15':length&0xffffffff,'%r14':SOURCE,'%r12':MSG,'%rdx':length&0xffffffff,'%rsi':SOURCE,'%rdi':MSG,'%rsp':WRAPPER_STACK};pc=start;trace=[];requests=[];sizes=[];flags=[]
            while True:
                t=run(dg,pc,list(calls)+[end],m,regs);m=t['memory'];regs=t['registers'];trace+=t['trace']
                if t['stop']==end:break
                i=calls[t['stop']]
                if '<MSG_WriteBits>' not in rows[i]['operands']:raise ValueError('unexpected data helper')
                request=[regs[k] for k in ['rdi','rsi','rdx']];requests.append(request)
                m,child_requests,error,child=writer(graphs,m,request[2],request[1]);trace+=child
                if child_requests or error:raise ValueError('unexpected OOB byte write outcome')
                sizes.append(get_word(m,MSG+28));flags.append(get_word(m,MSG+4));pc=rows[i+1]['address']
                if len(requests)>32:raise ValueError('data fixture bound')
            size=get_word(m,MSG+28)
            data.append({'length':length,'capacity':capacity,'initial_overflow':flag,'source_hex':bytes(source).hex(),'requests':requests,'post_write_sizes':sizes,'post_overflow_flags':flags,'size':size,'bit':get_word(m,MSG+36),'overflow_flag':get_word(m,MSG+4),'oldsize_word':get_word(m,old),'range_counter':get_word(m,over),'buffer_hex':bytes(m[DATA+i] for i in range(64)).hex(),'source_canaries':[m[SOURCE-1],m[SOURCE+32]],'buffer_canaries':[m[DATA-1],m[DATA+64]],**trace_fields({'trace':trace},False)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'scalar_oob':scalar,'reader_tails':tails,'data_oob':data}
    write_json(Path('docs/research/message-wrappers.json'),out)
    print('Recovered scalar OOB wrappers, supplied reader tails and data overflow compositions')


if __name__=='__main__':main()