Module refinery.lib.emulator.se

"""
Implements `refinery.lib.emulator.interface.Emulator` for the speakeasy backend.
"""
from __future__ import annotations

from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, TypeVar

from refinery.lib.emulator.abstract import (
    EmulationError,
    EmulationTimeout,
    Emulator,
    MemAccess,
    Register,
)
from refinery.lib.emulator.uc_shared import get_access_map
from refinery.lib.executable import ET, Arch
from refinery.lib.shared.speakeasy import speakeasy as se
from refinery.lib.vfs import VirtualFileSystem

if TYPE_CHECKING:
    from speakeasy import Speakeasy as Se
    from speakeasy.common import Hook as SeHook
    from speakeasy.memmgr import MemMap
else:
    class Se:
        pass


class SpeakeasyNotInitialized(EmulationError):
    def __init__(self) -> None:
        super().__init__('Speakeasy was unexpectedly not initialized.')


_T = TypeVar('_T')


class _InstructionBudget:
    """
    Bounds how many instructions Speakeasy emulates. Unicorn and Icicle bound a run with their
    engine's native instruction limit (the `count` argument of `emu_start` and the `icount_limit`
    register), but Speakeasy drives run termination through its own code hooks and exposes no
    detectable native limit, so a counting hook is the reliable and consistent choice here. The
    instance is installed as a per-instruction code hook: each call accounts for one instruction
    about to execute, and when the limit is reached it stops the engine and records that the budget
    was exhausted. Since `expired` is set only by this counter, the backend can tell a genuine
    timeout apart from any other stop reason and raises
    `refinery.lib.emulator.abstract.EmulationTimeout` accordingly.
    """
    __slots__ = ('limit', 'stop', 'count', 'expired')

    def __init__(self, limit: int, stop: Callable[[], Any]):
        self.limit = limit
        self.stop = stop
        self.count = 0
        self.expired = False

    def __call__(self, *args, **kwargs) -> bool:
        if self.count < self.limit:
            self.count += 1
        else:
            self.expired = True
            self.stop()
        return True


class SpeakeasyEmulator(Emulator[Se, str, _T]):
    """
    A Speakeasy-based emulator. Speakeasy only supports PE files, but it has support for several
    Windows API routines which can be an advantage.
    """

    speakeasy: Se

    def _init(self):
        self._regs: dict[str, Register[str]] = {}
        self._module: Any = None
        self._driver: Any = None

    class _singlestep:
        def __init__(self):
            self.stepped = False

        def __call__(self, se: Se, *_, **kw):
            if self.stepped:
                self.stepped = False
                se.stop()
            else:
                self.stepped = True
            return True

    def _reset(self):
        exe = self.exe
        if exe.type not in (ET.PE, ET.BLOB):
            raise NotImplementedError(F'Speakeasy cannot handle executables of type {exe.type.name}.')
        try:
            arch = {
                Arch.X32: 'x86',
                Arch.X64: 'x64',
            }[exe.arch()]
        except KeyError as KE:
            raise NotImplementedError(F'Speakeasy cannot handle executables of arch {exe.arch().name}') from KE

        emu = self.speakeasy = se.Speakeasy()
        self._driver = None

        with VirtualFileSystem() as vfs:
            db = bytes(exe.data)
            vf = vfs.new(db)
            if exe.blob:
                self._module = None
                self.base = emu.load_shellcode(vf.path, data=db, arch=arch)
            else:
                self._module = emu.load_module(vf.path, data=db)
                self.base = self._module.base

        if emu.emu is None:
            raise RuntimeError('emulator failed to initialize')

        self._end_hook_s = None
        self._end_hook_d = None

        self._single_step_hook_s = emu.add_code_hook(self._singlestep())
        self._single_step_hook_d = emu.add_dyn_code_hook(self._singlestep())
        self._disable_single_step()

        # Speakeasy's wall-clock run timeout (config.timeout) is a frozen pydantic field with no
        # public setter; copy the config to disable it. Termination is instead governed by the
        # return_hook, an explicit end address, or config.max_instructions.
        emu.emu.config = emu.emu.config.model_copy(update={'timeout': 0})

        # Speakeasy registers one native dispatch hook per event that stops at the first callback
        # returning False, with its builtin _hook_mem_unmapped added last; that builtin clobbers
        # our memory hooks (and single-stepping). 2.0.0b3 has no way to make consumer hooks coexist
        # with the builtins, so we suppress them and replace the interrupt handler ourselves.
        emu.emu.add_interrupt_hook(cb=emu.emu._hook_interrupt)
        emu.emu.builtin_hooks_set = True

        # The suppressed builtin was also what detected the run-ending fetch into the unmapped
        # return_hook, so we detect it ourselves (_uc_hook_return) to stop a run when it returns.
        emu.add_mem_invalid_hook(self._uc_hook_return)

        self._access_map = get_access_map()

        if self.hooks.CodeExecute:
            emu.add_code_hook(self._uc_hook_code)
            emu.add_dyn_code_hook(self._uc_hook_code)

        if self.hooks.MemoryRead:
            emu.add_mem_read_hook(self._uc_hook_mem_read)

        if self.hooks.MemoryWrite:
            emu.add_mem_write_hook(self._uc_hook_mem_write)

        if self.hooks.MemoryError:
            emu.add_mem_invalid_hook(self._uc_hook_mem_error)

        if self.hooks.ApiCall:
            emu.add_api_hook(self.hook_api_call, '*', '*')

    def _uc_hook_code(self, emu: Se, address: int, size: int, *_) -> bool:
        return self.hook_code_execute(emu, address, size, self.state)

    def _uc_hook_mem_read(self, emu: Se, access: int, address: int, size: int, value: int, state: _T | None = None) -> bool:
        access = self._access_map.get(access, MemAccess.Unknown)
        return self.hook_mem_read(emu, access, address, size, value, state)

    def _uc_hook_mem_write(self, emu: Se, access: int, address: int, size: int, value: int, state: _T | None = None) -> bool:
        access = self._access_map.get(access, MemAccess.Unknown)
        return self.hook_mem_write(emu, access, address, size, value, state)

    def _uc_hook_mem_error(self, emu: Se, access: int, address: int, size: int, value: int, state: _T | None = None) -> bool:
        access = self._access_map.get(access, MemAccess.Unknown)
        return self.hook_mem_error(emu, access, address, size, value, state)

    def _uc_hook_return(self, emu: Se, access: int, address: int, size: int, value: int, state: _T | None = None):
        inner = self.speakeasy.emu
        if inner is not None and address in (inner.return_hook, inner.exit_hook):
            inner.stop()
            return True
        return None

    def _enable_single_step(self):
        hd = self._single_step_hook_d
        hs = self._single_step_hook_s
        if hd is None or hs is None:
            raise RuntimeError('single stepping hooks failed to be installed')
        hd.cb.stepped = False
        hs.cb.stepped = False
        hd.enable()
        hs.enable()

    def _disable_single_step(self):
        if hook := self._single_step_hook_d:
            hook.disable()
        if hook := self._single_step_hook_s:
            hook.disable()

    @property
    def stack_region(self):
        emu = self.speakeasy
        tos = self.sp
        mms: list[MemMap] = emu.get_mem_maps()
        if tos != emu.get_stack_ptr():
            raise EmulationError('Unexpected stack pointer misalignment')
        try:
            sm, = (mm for mm in mms if tos in range(mm.base, mm.base + mm.size))
        except Exception:
            raise EmulationError('Ambiguous memory, unable to locate the stack.')
        return sm

    def _map_update(self):
        self._memorymap.clear()
        if self.speakeasy.emu is None:
            raise SpeakeasyNotInitialized
        for mm in self.speakeasy.get_mem_maps():
            self._memorymap.addi(mm.base, mm.size)

    def malloc(self, size: int) -> int:
        return self.speakeasy.mem_alloc(size)

    def push(self, val: int, size: int | None = None):
        if size is None:
            size = self.exe.pointer_size_in_bytes
        easy = self.speakeasy
        sp = easy.get_stack_ptr()
        bv = val.to_bytes(size, self.exe.byte_order().value)
        sp -= size
        easy.mem_write(sp, bv)
        easy.set_stack_ptr(sp)

    def morestack(self):
        # Speakeasy manages its own stack; grow it downward by one page when asked. push() is
        # overridden below and does not need this, so it is only a best-effort safety net.
        region = self.stack_region
        self.map(region.base - self.alloc_size, self.alloc_size)

    class _stop:
        hook: SeHook | None
        address: int | None

        def __init__(self, address: int | None = None):
            self.address = address
            self.hook = None

        def __call__(self, spky: Se, address: int | None = None, *_):
            if hook := self.hook:
                if address == self.address:
                    spky.stop()
                    hook.disable()

    _end_hook_s: _stop | None
    _end_hook_d: _stop | None

    def _remove_hook(self, hook: SeHook | None):
        if hook is None:
            return
        try:
            hook.emu_eng.hook_remove(hook.handle)
        except KeyError:
            # Speakeasy installs its Unicorn hooks by calling uc_hook_add directly, bypassing the
            # binding's own callback table; uc_hook_del then unhooks the native hook successfully
            # but the binding's redundant bookkeeping delete raises KeyError. Speakeasy's own
            # close() swallows this the same way; the native hook is already gone at this point.
            pass
        emu = self.speakeasy.emu
        assert emu is not None
        for hooklist in emu.hooks.values():
            assert isinstance(hooklist, list)
            for k, h in enumerate(hooklist):
                if h is hook:
                    del hooklist[k]
                    break

    def _set_end(self, end: int | None):
        if h := self._end_hook_s:
            self._remove_hook(h.hook)
        if h := self._end_hook_d:
            self._remove_hook(h.hook)
        if end is None:
            self._end_hook_s = None
            self._end_hook_d = None
        else:
            self._end_hook_s = h = self._stop(end)
            h.hook = self.speakeasy.add_code_hook(h, end, end + 1)
            self._end_hook_d = h = self._stop(end)
            h.hook = self.speakeasy.add_dyn_code_hook(h)

    def _ensure_context(self) -> list:
        """
        Idempotently establish the process (Win32) or driver (kernel) context required to run the
        loaded module at an arbitrary address through the public `Speakeasy.call`, and return the
        argument list for its entry. The facade exposes no way to set up a run context without
        also executing the module's entry point, so this reaches into Speakeasy 2.0.0b3 internals.
        """
        inner = self.speakeasy.emu
        assert inner is not None
        module = self._module
        if isinstance(inner, se.WinKernelEmulator):
            if self._driver is None:
                regdefs = se.winenv.defs.registry.reg
                drv = self._driver = inner.create_driver_object(pe=module)
                svc = inner.regman.create_key(drv.reg_path)
                svc.create_value('ImagePath', regdefs.REG_EXPAND_SZ, module.emu_path)
                svc.create_value('Type', regdefs.REG_DWORD, 0x1)
                svc.create_value('Start', regdefs.REG_DWORD, 0x3)
                svc.create_value('ErrorControl', regdefs.REG_DWORD, 0x1)
                inner.regman.create_key(drv.reg_path + '\\Parameters')
            return [self._driver.address, self._driver.reg_path_ptr]
        if not inner.processes:
            process = se.windows.objman.Process(
                inner,
                path=module.emu_path,
                base=module.base,
                pe=module,
                cmdline=inner.command_line,
            )
            inner.curr_process = process
            inner.om.objects.update({process.address: process})
            if mm := inner.get_address_map(module.base):
                mm: MemMap
                mm.process = process
            inner.alloc_peb(process)
        return []

    @contextmanager
    def _preserve_cpu_state(self, start: int):
        """
        Speakeasy resets the stack pointer and the argument registers (via reset_stack and
        set_func_args) when it kicks off a run, discarding whatever CPU state the caller set up
        before emulate(). Snapshot the general purpose registers and stack pointer and reinstate
        them with a one-shot code hook at the run's first instruction, so a shellcode run continues
        from the current CPU state like the unicorn and icicle backends do.
        """
        registers = [reg.code for reg in self.general_purpose_registers()]
        registers.append(self._reg_sp)
        snapshot = {reg: self._get_register(reg) for reg in registers}
        restored = False

        def restore(spky: Se, address: int, size: int, *_):
            nonlocal restored
            if not restored and address == start:
                restored = True
                for reg, value in snapshot.items():
                    self._set_register(reg, value)
            return True

        hook = self.speakeasy.add_code_hook(restore, start, start + 1)
        try:
            yield
        finally:
            self._remove_hook(hook)

    def _emulate(self, start: int, end: int | None = None, timeout: int | None = None):
        spk = self.speakeasy

        if (inner := spk.emu) is None:
            raise SpeakeasyNotInitialized

        self._set_end(end)

        budget = None
        countdown_hooks = []
        if timeout is not None:
            budget = _InstructionBudget(timeout, spk.stop)
            countdown_hooks.append(spk.add_code_hook(budget))
            countdown_hooks.append(spk.add_dyn_code_hook(budget))

        try:
            if inner.get_current_run():
                spk.resume(start)
            elif self.exe.blob:
                offset = start - self.base
                if offset < 0:
                    raise ValueError(F'invalid offset 0x{start:X} specified; base address is 0x{self.base:X}')
                with self._preserve_cpu_state(start):
                    spk.run_shellcode(self.base, offset=offset)
            else:
                spk.call(start, self._ensure_context())
        finally:
            for hook in countdown_hooks:
                self._remove_hook(hook)

        if budget is not None and budget.expired:
            raise EmulationTimeout(budget.count)

    def halt(self):
        return self.speakeasy.stop()

    def _set_register(self, register: str, v: int):
        return self.speakeasy.reg_write(register, v)

    def _get_register(self, register: str) -> int:
        return self.speakeasy.reg_read(register)

    def _lookup_register(self, var: str) -> Register[str]:
        try:
            reg = self._regs[var]
        except KeyError:
            try:
                size = self.measure_register_size(var)
            except Exception:
                raise LookupError(var)
            else:
                reg = self._regs[var] = Register(var, var, size)
        return reg

    def _map(self, address: int, size: int):
        spksy = self.speakeasy
        if (emu := spksy.emu) is None:
            raise SpeakeasyNotInitialized
        if emu.get_address_map(address):
            raise ValueError(address)
        if mm := emu.get_reserve_map(address):
            mm: MemMap = emu.get_address_map(emu.mem_map_reserve(mm.base))
            if address not in range(mm.base, mm.base + mm.size):
                raise RuntimeError(F'Speakeasy claimed to map 0x{address:X} in map 0x{mm.base:X}-0x{mm.base + mm.size:X}.')
            map_size = mm.size
            map_base = mm.base
            _new_size = size - map_size + address - map_base
            _new_base = map_base + map_size
            if _new_size > 0 and self._map(_new_base, _new_size) != _new_base:
                raise RuntimeError(F'Attempting to remain rest of size 0x{_new_size:X} at 0x{_new_base:X} failed.')
            return address
        else:
            alloc = spksy.mem_alloc(size, address)
            if alloc != address:
                spksy.mem_free(alloc)
                raise LookupError(F'Unable to allocate {size} bytes at address 0x{address:X} because Speakeasy has reserved this region.')
            return alloc

    def _mem_write(self, address: int, data: bytes):
        return self.speakeasy.mem_write(address, data)

    def _mem_read(self, address: int, size: int):
        return self.speakeasy.mem_read(address, size)