#!/usr/bin/env python3
"""Byte-guarded automatic CFG lift of Factorio's original MapHeader constructor."""
import argparse
import hashlib
import json
import struct
import subprocess
from pathlib import Path
from a64_program import ProgramCompiler
from lift_rng import arm64_image, file_range

ROOT = Path(__file__).resolve().parents[1]
ADDRESS, SIZE = 0x1014C6AAC, 332
EXPECTED = '4f940fe2feb23b8301d080d3678634fa33891d4d53de87c844d70c1319282ff6'
READ = 0x10035E264
ENTRY, ENTRY_SIZE = 0x10010582C, 4
ENTRY_HASH = 'cb0ac2c0409e9c0c1a8de8acf7ce68abc43ca0ab4fb20c2a97c27daab156df81'


def run(*args):
    subprocess.run(list(map(str, args)), cwd=ROOT, check=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('binary', type=Path)
    args = parser.parse_args()
    image, slice_offset = arm64_image(args.binary.read_bytes())
    offset = file_range(image, ADDRESS, SIZE)
    body = image[offset:offset+SIZE]
    if hashlib.sha256(body).hexdigest() != EXPECTED:
        raise ValueError('Unrecognized MapHeader constructor bytes')
    entry_offset = file_range(image, ENTRY, ENTRY_SIZE)
    entry_body = image[entry_offset:entry_offset+ENTRY_SIZE]
    if hashlib.sha256(entry_body).hexdigest() != ENTRY_HASH:
        raise ValueError('Unrecognized MapHeader public constructor entry')
    compiler = ProgramCompiler({READ: 'read_or_throw'})
    wat = compiler.compile({ENTRY:entry_body,ADDRESS:body}, ENTRY).replace('(import "env"', '(import "map_host"')
    out = ROOT / 'build/map-header'
    out.mkdir(parents=True, exist_ok=True)
    (out / 'factorio_map_header.wat').write_text(wat)
    manifest = {'symbol': 'MapHeader::MapHeader(MapDeserialiser&)', 'address': hex(ADDRESS),
        'bytes': SIZE, 'sha256': EXPECTED, 'source_file_offset': slice_offset+offset,
        'instructions': compiler.instructions, 'host_calls': compiler.calls,
        'rea_bytes_evidence': 'ev_edc9090e494bf4c92f5c0e5d9f4366272062c6d2d3d758442448ab720bf4e422',
        'public_entry': {'address':hex(ENTRY),'bytes':ENTRY_SIZE,'sha256':ENTRY_HASH,
                         'rea_bytes_evidence':'ev_525aeb761646ca31af2fbacfc51ec786e6036b6259fd09505063144d0a5feba1'},
        'scope': 'version-dependent header constructor; declared memory-stream adapter; no C++ unwinding'}
    (out / 'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
    assembly = ['.text', '.p2align 2', '.globl _factorio_map_header_original', '_factorio_map_header_original:',
                'b Lmap_header_c2', 'Lmap_header_c2:'] # Relocate the C1 tail B to the copied C2 body.
    for (word,) in struct.iter_unpack('<I', body):
        assembly.append('bl _native_map_header_read' if word & 0xfc000000 == 0x94000000 else f'.long 0x{word:08x}')
    (out / 'oracle.S').write_text('\n'.join(assembly)+'\n')
    run(ROOT / 'build/wabt/wat2wasm', out / 'factorio_map_header.wat', '-o', out / 'factorio_map_header.wasm')
    run(ROOT / 'build/wabt/wasm2c', out / 'factorio_map_header.wasm', '-o', out / 'factorio_map_header.c', '--module-name=factorio_map_header')
    runtime = ROOT / 'third_party/wabt/wasm2c'
    includes = ['-I'+str(p) for p in (out, ROOT/'runtime', runtime)]
    flags = ['-O2', '-DWASM_RT_USE_MMAP=0', '-DWASM_RT_MEMCHECK_BOUNDS_CHECK=1']
    sources = [out/'factorio_map_header.c', ROOT/'runtime/map_header_bridge.c']
    run('xcrun', 'clang', *flags, *includes, *sources, runtime/'wasm-rt-impl.c', runtime/'wasm-rt-mem-impl.c',
        out/'oracle.S', ROOT/'tests/map_header_equivalence.c', '-o', out/'verify_map_header')
    run(out/'verify_map_header')
    sdk = subprocess.check_output(['xcrun', '--sdk', 'iphoneos', '--show-sdk-path'], text=True).strip()
    objects = []
    for source in sources:
        obj = out/(source.stem+'.ios.o')
        run('xcrun', '--sdk', 'iphoneos', 'clang', '-target', 'arm64-apple-ios15.0', '-isysroot', sdk,
            *flags, *includes, '-c', source, '-o', obj)
        objects.append(obj)
    run('xcrun', 'libtool', '-static', '-o', out/'libfactorio_map_header_ios.a', *objects)
    (out/'verification.json').write_text(json.dumps({'status':'pass', 'native_comparisons':20000,
        'instructions':len(compiler.instructions), 'host_call_sites':len(compiler.calls),
        'whole_map_loading':False, 'scope':'controlled payload at constructor stream position'}, indent=2)+'\n')
    print('PASS: original MapHeader CFG lift, native oracle comparisons, and iOS static library')


if __name__ == '__main__':
    main()