#!/usr/bin/env python3
"""Writer guard slices and native byte-codec compositions with modeled trees."""
import copy
import gzip
import hashlib
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 evaluate,get_word,put_word
from analyze_engine_control import control_flow
from analyze_huffman_encoding import send
from analyze_qvm import write_json
from analyze_vm_extensions import native_functions,verified

MSG,DATA,STACK=0x90000,0xa0000,0xb0000


def run(g,start,stops,m,regs):return evaluate(g,start,stops,m,regs,limit=10000,capture_registers=True)


def install_tree(m,s,base):
    original=0x10000;delta=base-original
    def pointer(v):return v+delta if v else 0
    for i in range(0xc830):m[base+i]=0
    put_word(m,base,s['node_count']);put_word(m,base+4,s['pointer_cell_count'])
    for offset,key in [(8,'root'),(16,'list_start'),(24,'tail_slot'),(0x828,'free_head')]:put_word(m,base+offset,pointer(s[key]),8)
    for symbol,value in s['lookup']:put_word(m,base+32+8*symbol,pointer(value),8)
    for i,value in enumerate(s['pointer_cells']):put_word(m,base+0xb030+8*i,pointer(value),8)
    for node in s['nodes']:
        address=pointer(node['address'])
        for offset,key in [(0,'left'),(8,'right'),(16,'parent'),(24,'next'),(32,'previous'),(40,'head_cell')]:put_word(m,address+offset,pointer(node[key]),8)
        put_word(m,address+48,node['weight']);put_word(m,address+52,node['symbol'])
    return pointer(s['root'])


def offset_transmit(g,sg,m,regs):
    graph=copy.deepcopy(g);rows=graph['instructions']
    for x in rows:
        if x['opcode']=='pushq':x['opcode']='nop';x['operands']=''
    idx=next(i for i,x in enumerate(rows) if x['opcode']=='callq');end=next(x['address'] for x in rows if x['opcode']=='popq')
    t=run(graph,rows[0]['address'],[rows[idx]['address'],end],m,regs);trace=t['trace'];m=t['memory']
    if t['stop']!=end:
        r=t['registers'];child=send(sg,m,r['rdi'],r['rsi'],r['rdx']);m=child['memory'];trace+=child['trace']
        t=run(graph,rows[idx+1]['address'],[end],m,r);m=t['memory'];trace+=t['trace']
    return m,trace


def byte_write(graphs,m,value):
    g=graphs['MSG_WriteBits'];rows=g['instructions'];start=next(x['address'] for x in rows if x['opcode']=='addl' and '(%rip)' in x['operands']);end=next(x['address'] for x in rows if x['opcode']=='movq' and x['operands'].split('#')[0].strip()=='0x8(%rsp), %rbx');calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'}
    regs={'%rdi':MSG,'%rbp':MSG,'%rbx':value,'%rdx':8,'%rsp':STACK};trace=[];requests=[]
    while True:
        t=run(g,start,list(calls)+[end],m,regs);trace+=t['trace'];m=t['memory'];regs=t['registers']
        if t['stop']==end:break
        idx=calls[t['stop']]
        if '<Huff_offsetTransmit>' not in rows[idx]['operands']:raise ValueError('unexpected byte writer helper')
        requests.append([regs[k] for k in ['rdi','rsi','rdx','rcx']]);m,child=offset_transmit(graphs['Huff_offsetTransmit'],graphs['send'],m,regs);trace+=child;start=rows[idx+1]['address']
    return m,trace,requests


def byte_read(graphs,m):
    g=graphs['MSG_ReadBits'];rows=g['instructions'];start=site(g,'testl','%esi, %esi');end=next(x['address'] for x in rows if x['opcode']=='movl' and x['operands'].split('#')[0].strip()=='%ebp, %eax');calls={x['address']:i for i,x in enumerate(rows) if x['opcode']=='callq'}
    regs={'%r12':MSG,'%r15':8,'%rsi':8,'%rsp':STACK};trace=[]
    while True:
        t=run(g,start,list(calls)+[end],m,regs);trace+=t['trace'];m=t['memory'];regs=t['registers']
        if t['stop']==end:break
        idx=calls[t['stop']]
        if '<Huff_offsetReceive>' not in rows[idx]['operands']:raise ValueError('unexpected byte reader helper')
        h=graphs['Huff_offsetReceive'];child=run(h,h['instructions'][0]['address'],[x['address'] for x in h['instructions'] if x['opcode']=='retq'],m,regs);m=child['memory'];trace+=child['trace'];start=rows[idx+1]['address']
    return m,trace,regs['rbp']&0xffffffff


def main():
    cp=Path('docs/research/huffman-codebook.json');code=json.loads(cp.read_text());initial=json.loads(Path('docs/research/huffman-initialization.json').read_text());out={'native':{},'codebook_report_sha256':hashlib.sha256(cp.read_bytes()).hexdigest(),'limitations':'Linux ELF64 synthetic native writer/reader/helper compositions using MODEL-DERIVED supplied seed trees. No original process/native callee/runtime/codebook capture or real packet executes. Saves/restores/PUSH/POP excluded; explicit caller restoration. Byte-width only; general raw/signed/multi-byte message codecs, capacities, truncation, packet/snapshot/all-platform and full native seed equivalence remain open.'}
    sequences=[[0],[255],[0,1,127,128,255],list(range(256)),[0,255,128,1]*8,[256,257,0xffffffff,0x1234]]
    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());graphs={n:control_flow(fs[n]) for n in ['MSG_WriteBits','MSG_ReadBits','Huff_offsetTransmit','Huff_offsetReceive','send']}
        if code['native'][label]['source']!=r['source'] or code['native'][label]['disassembly_sha256']!=a['uncompressed_sha256']:raise ValueError('codebook provenance')
        g=graphs['MSG_WriteBits'];rows=g['instructions'];start=next(x['address'] for x in rows if x['opcode']=='addl' and '(%rip)' in x['operands']);end=next(x['address'] for x in rows if x['opcode']=='movq' and x['operands'].split('#')[0].strip()=='0x8(%rsp), %rbx');mode=site(g,'movl','0x8(%rbp), %r9d');error=next(x['address'] for x in rows if '<Com_Error>' in x['operands']);old=global_address(next(x for x in rows if x['opcode']=='addl' and '(%rip)' in x['operands']));over=global_address(next(x for x in rows if x['opcode']=='addl' and '$0x1,' in x['operands'] and '(%rip)' in x['operands']));guards=[]
        for width,value,remaining in itertools.product([0,1,7,8,16,31,32,33,-1,-8,-31,-32],[0,127,128,255,256,0xffffffff,0x7fffffff,0x80000000],[0,3,4]):
            m={};put_word(m,MSG+24,100+remaining);put_word(m,MSG+28,100);put_word(m,MSG+4,0);put_word(m,old,7);put_word(m,over,11)
            t=run(g,start,[end,mode,error],m,{'%rdi':MSG,'%rbp':MSG,'%rbx':value,'%rdx':width&0xffffffff,'%rsp':STACK});regs=t['registers']
            guards.append({'requested_width':width,'input_word':value,'remaining_bytes':remaining,'outcome':'space_overflow' if t['stop']==end else 'width_error' if t['stop']==error else 'mode_dispatch','normalized_width_word':regs['r14'],'overflow_flag':get_word(t['memory'],MSG+4),'range_overflow_counter':get_word(t['memory'],over),'oldsize_word':get_word(t['memory'],old),**trace_fields(t,False)})
        streams=[];base=initial['native'][label]['entry_request']['huff_init_pointer'];s=code['model_derived_seed_training']['final_state']
        for sequence,bit in itertools.product(sequences,[0,7,8]):
            m={};install_tree(m,s,base);install_tree(m,s,base+0xc830)
            for x in graphs['MSG_ReadBits']['instructions']:
                if x['opcode']=='movq' and '(%rip)' in x['operands']:put_word(m,global_address(x),s['root']+base+0xc830-0x10000,8)
            for i in range(4096):m[DATA+i]=0
            m[DATA-1]=m[DATA+4096]=0x7b
            for offset,value in [(4,0),(8,0),(24,4096),(28,bit//8+1),(32,0),(36,bit)]:put_word(m,MSG+offset,value)
            put_word(m,MSG+16,DATA,8);put_word(m,old,0);put_word(m,over,0);trace=[];writes=[];requests=[]
            for value in sequence:
                m,part,req=byte_write(graphs,m,value);trace+=part;requests+=req;writes.append([get_word(m,MSG+36),get_word(m,MSG+28)])
            size=get_word(m,MSG+28);payload=[m[DATA+i] for i in range(size)];put_word(m,MSG+36,bit);put_word(m,MSG+32,0);decoded=[];reads=[]
            for _ in sequence:
                m,part,value=byte_read(graphs,m);trace+=part;decoded.append(value);reads.append([get_word(m,MSG+36),get_word(m,MSG+32)])
            streams.append({'symbols':sequence,'initial_bit':bit,'payload_bytes':payload,'write_bit_size_pairs':writes,'read_bit_count_pairs':reads,'decoded_words':decoded,'transmit_requests':requests,'overflow_flag':get_word(m,MSG+4),'range_overflow_counter':get_word(m,over),'oldsize_word':get_word(m,old),'canaries':[m[DATA-1],m[DATA+4096]],**trace_fields({'trace':trace},False)})
        out['native'][label]={'source':r['source'],'disassembly_sha256':a['uncompressed_sha256'],'graphs':graphs,'writer_guard_probes':guards,'byte_streams':streams}
    write_json(Path('docs/research/message-byte-codec.json'),out);print('Recovered 576 writer guards and 36 modeled-tree native message-byte streams')


if __name__=='__main__':main()