#!/usr/bin/env python3
"""
emulate.py — run the Linux static-ELF assembly builds (i386, x86-64, riscv64) on any host via the Unicorn CPU emulator.

    pip install unicorn                                    # (Unicorn 2.x; a venv is fine)
    python3 emulate.py <static-elf> [args...]              # stdin/stdout/exit code are passed through

Why: this repo is developed on Apple silicon, which cannot execute 32-bit x86 or RISC-V code at all. The
assembly programs are cross-assembled with `zig cc -target {x86,x86_64,riscv64}-linux-musl -nostdlib -static`
and executed here for real by emulating the CPU; only the handful of Linux syscalls they use are serviced by
this harness:  read(0) · write(1) · getrandom → os.urandom (§5.2, the host OS CSPRNG) · exit/exit_group.

Loader: PT_LOAD segments mapped at their vaddr; a 1 MiB stack with the SysV process-entry layout
(argc, argv[], NULL, envp NULL, auxv AT_NULL); PC = e_entry. Nothing else — no dynamic loading, no TLS.
"""
import os, struct, sys
from unicorn import Uc, UC_ARCH_X86, UC_ARCH_RISCV, UC_MODE_32, UC_MODE_64, UC_MODE_RISCV64, UC_HOOK_INTR, UC_HOOK_INSN, UcError
from unicorn.x86_const import (UC_X86_REG_EAX, UC_X86_REG_EBX, UC_X86_REG_ECX, UC_X86_REG_EDX, UC_X86_REG_ESP,
                               UC_X86_REG_RAX, UC_X86_REG_RDI, UC_X86_REG_RSI, UC_X86_REG_RDX, UC_X86_REG_RSP,
                               UC_X86_INS_SYSCALL)
from unicorn.riscv_const import UC_RISCV_REG_A0, UC_RISCV_REG_A1, UC_RISCV_REG_A2, UC_RISCV_REG_A7, UC_RISCV_REG_SP

PAGE = 0x1000
def align_down(x): return x & ~(PAGE - 1)
def align_up(x): return (x + PAGE - 1) & ~(PAGE - 1)

# syscall numbers per ABI: (read, write, getrandom, exit, exit_group)
ABI = {
    'i386':    dict(read=3,  write=4,  getrandom=355, exit=1,  exit_group=252),
    'x86_64':  dict(read=0,  write=1,  getrandom=318, exit=60, exit_group=231),
    'riscv64': dict(read=63, write=64, getrandom=278, exit=93, exit_group=94),
}

def load_elf(path):
    data = open(path, 'rb').read()
    assert data[:4] == b'\x7fELF', 'not an ELF file'
    bits = 32 if data[4] == 1 else 64
    machine = struct.unpack_from('<H', data, 18)[0]
    arch = {3: 'i386', 62: 'x86_64', 243: 'riscv64'}.get(machine)
    if arch is None: sys.exit(f'unsupported e_machine {machine}')
    if bits == 32:
        e_entry, e_phoff = struct.unpack_from('<II', data, 24)
        e_phentsize, e_phnum = struct.unpack_from('<HH', data, 42)
    else:
        e_entry, e_phoff = struct.unpack_from('<QQ', data, 24)
        e_phentsize, e_phnum = struct.unpack_from('<HH', data, 54)
    segs = []
    for i in range(e_phnum):
        off = e_phoff + i * e_phentsize
        if bits == 32:
            p_type, p_offset, p_vaddr, _, p_filesz, p_memsz, _, _ = struct.unpack_from('<IIIIIIII', data, off)
        else:
            p_type, _, p_offset, p_vaddr, _, p_filesz, p_memsz, _ = struct.unpack_from('<IIQQQQQQ', data, off)
        if p_type == 1:  # PT_LOAD
            segs.append((p_vaddr, data[p_offset:p_offset + p_filesz], p_memsz))
    return arch, bits, e_entry, segs

def main():
    if len(sys.argv) < 2:
        print(__doc__); return 2
    path, args = sys.argv[1], sys.argv[1:]          # argv[0] = program path, like a real exec
    arch, bits, entry, segs = load_elf(path)
    uc = {'i386': lambda: Uc(UC_ARCH_X86, UC_MODE_32), 'x86_64': lambda: Uc(UC_ARCH_X86, UC_MODE_64),
          'riscv64': lambda: Uc(UC_ARCH_RISCV, UC_MODE_RISCV64)}[arch]()
    sysno = ABI[arch]

    # map segments
    mapped = []
    for vaddr, filedata, memsz in segs:
        lo, hi = align_down(vaddr), align_up(vaddr + memsz)
        for (a, b) in mapped:                       # merge overlaps crudely: skip already-mapped pages
            if lo < b and hi > a: lo, hi = min(lo, a), max(hi, b)
        try: uc.mem_map(lo, hi - lo)
        except UcError: pass                        # already mapped by an overlapping segment
        mapped.append((lo, hi))
        uc.mem_write(vaddr, filedata)

    # stack + SysV entry layout
    stack_top = 0xbfff0000 if bits == 32 else 0x7ffff0000000
    stack_size = 1 << 20
    uc.mem_map(stack_top - stack_size, stack_size)
    ptr = 8 if bits == 64 else 4
    fmt = '<Q' if bits == 64 else '<I'
    sp = stack_top
    argv_ptrs = []
    for a in reversed(args):
        b = a.encode() + b'\0'
        sp -= len(b); uc.mem_write(sp, b); argv_ptrs.append(sp)
    argv_ptrs.reverse()
    sp &= ~0xF
    words = [len(args)] + argv_ptrs + [0] + [0] + [0, 0]      # argc, argv..., NULL, envp NULL, auxv AT_NULL
    sp -= ptr * len(words)
    sp &= ~0xF
    for i, w in enumerate(words): uc.mem_write(sp + i * ptr, struct.pack(fmt, w))
    if arch == 'i386': uc.reg_write(UC_X86_REG_ESP, sp)
    elif arch == 'x86_64': uc.reg_write(UC_X86_REG_RSP, sp)
    else: uc.reg_write(UC_RISCV_REG_SP, sp)

    stdin_buf = sys.stdin.buffer
    out = sys.stdout.buffer
    exit_code = [None]

    def do_syscall(no, a0, a1, a2):
        if no == sysno['write']:
            if a0 in (1, 2): (out if a0 == 1 else sys.stderr.buffer).write(bytes(uc.mem_read(a1, a2))); return a2
            return -9                                       # EBADF
        if no == sysno['read']:
            if a0 != 0: return -9
            data = stdin_buf.read(a2)
            uc.mem_write(a1, data); return len(data)
        if no == sysno['getrandom']:
            uc.mem_write(a0, os.urandom(a1)); return a1     # §5.2 — host OS CSPRNG
        if no in (sysno['exit'], sysno['exit_group']):
            exit_code[0] = a0 & 0xFF; uc.emu_stop(); return 0
        sys.stderr.write(f'emulate.py: unsupported {arch} syscall {no}\n'); exit_code[0] = 127; uc.emu_stop(); return -38

    if arch == 'i386':
        def on_intr(uc, intno, _):
            if intno != 0x80: return
            r = do_syscall(uc.reg_read(UC_X86_REG_EAX), uc.reg_read(UC_X86_REG_EBX), uc.reg_read(UC_X86_REG_ECX), uc.reg_read(UC_X86_REG_EDX))
            uc.reg_write(UC_X86_REG_EAX, r & 0xFFFFFFFF)
        uc.hook_add(UC_HOOK_INTR, on_intr)
    elif arch == 'x86_64':
        def on_syscall(uc, _):
            r = do_syscall(uc.reg_read(UC_X86_REG_RAX), uc.reg_read(UC_X86_REG_RDI), uc.reg_read(UC_X86_REG_RSI), uc.reg_read(UC_X86_REG_RDX))
            uc.reg_write(UC_X86_REG_RAX, r & 0xFFFFFFFFFFFFFFFF)
        uc.hook_add(UC_HOOK_INSN, on_syscall, None, 1, 0, UC_X86_INS_SYSCALL)
    else:
        def on_ecall(uc, intno, _):
            if intno != 8: return                           # 8 = environment call from U-mode
            r = do_syscall(uc.reg_read(UC_RISCV_REG_A7), uc.reg_read(UC_RISCV_REG_A0), uc.reg_read(UC_RISCV_REG_A1), uc.reg_read(UC_RISCV_REG_A2))
            uc.reg_write(UC_RISCV_REG_A0, r & 0xFFFFFFFFFFFFFFFF)
            # (Unicorn resumes after the ecall by itself; do not bump PC)
        uc.hook_add(UC_HOOK_INTR, on_ecall)

    try:
        uc.emu_start(entry, 0)                             # runs until exit() stops it
    except UcError as e:
        if exit_code[0] is None:
            sys.stderr.write(f'emulate.py: {arch} fault: {e}\n'); return 139
    out.flush()
    return exit_code[0] if exit_code[0] is not None else 0

if __name__ == '__main__':
    sys.exit(main())