pcode_engine.py

from __future__ import annotations
from dataclasses import dataclass
from functools import partial

import pypcode
import ctypes

from engine_types import (
    Arg,
    BinaryOp,
    CallSite,
    ConditionalExpression,
    ConditionalSite,
    LoopsDict,
    Register,
    MemoryAccess,
    MemoryAccessType,
    UnaryOp,
)
from typing import Any, Callable, Optional
from frozendict import frozendict
from binary_function import BinaryFunction, FunctionBlock
from project import Project


class InstructionState:
    def __init__(self):
        self.regs: dict[int, Any] = {}
        self.unique: dict[int, Any] = {}
        self.ram: dict[int, Any] = {}
        self.stack: dict[int, Any] = {}
        self.last_callsite: Optional[CallSite] = None
        self.goto_state: dict[int, InstructionState] = dict()
        self.used_arguments: set[int] = set()

    def copy(self) -> InstructionState:
        new_state = InstructionState()
        new_state.regs = self.regs.copy()
        new_state.unique = self.unique.copy()
        new_state.ram = self.ram.copy()
        new_state.stack = self.stack.copy()
        new_state.last_callsite = self.last_callsite
        new_state.used_arguments = self.used_arguments
        return new_state


class Engine:

    def __init__(self, start: int, code: bytes, project: Project):
        self.project = project
        self.bin_func = BinaryFunction(start, code, project)
        self.instructions_state: dict[int, InstructionState] = dict()
        self.current_inst: int = 0
        self.previous_marks: list[int] = list()
        self._handlers: dict[pypcode.OpCode, Callable[[pypcode.PcodeOp], None]] = {}
        self.__put_handlers: dict[str, Callable[[int, Any], None]] = {}
        self.__get_handlers: dict[str, Callable[[int], Any]] = {}
        self.callsites: list[CallSite] = []
        self.conditional_sites: list[ConditionalSite] = []
        self.memory_accesses: list[MemoryAccess] = []
        self.addr_to_codeflow_conditional_site: dict[int, ConditionalSite] = {}
        self.current_blk: FunctionBlock = FunctionBlock(0, self.bin_func)  # Temporal initialization
        self.__unfinished_condsite: Optional[_UnfinishedConditionalSite] = None
        self.__conditional_move_condition: Optional[ConditionalSite] = None
        self._first_stack_access: Optional[Register] = None
        self._return_values: Optional[set[Any]] = None

        self._init_handlers()

        for current_addr in self.bin_func.code_flow_graph.traverse():
            blk = self.bin_func.blocks_dict_start_address[current_addr]
            parents = self.bin_func.code_flow_graph.get_parnets(current_addr)
            self.previous_marks = [self.bin_func.blocks_dict[parent].last_instruction_addr for parent in parents]
            self._analyze_block(blk)

    @property
    def loops_dict_start_address(self) -> LoopsDict:
        return self.bin_func.loops_dict_start_address

    @property
    def loops_dict(self) -> LoopsDict:
        return self.bin_func.loops_dict

    @property
    def return_values(self) -> set[Any]:
        if self._return_values is not None:
            return self._return_values

        self._return_values = set()

        callsites_by_addr = {cs.addr: cs for cs in self.callsites}

        for blk in self.bin_func.return_blocks.values():
            if blk.last_instruction_addr in callsites_by_addr:
                self._return_values.add(callsites_by_addr[blk.last_instruction_addr])
                continue

            ret_val = self.instructions_state[blk.last_instruction_addr].regs.get(self.project.arch_regs.ret)
            if ret_val is None:
                continue

            if isinstance(ret_val, ConditionalExpression):
                self._return_values.update(ret_val.collect_values())
            else:
                self._return_values.add(ret_val)

        return self._return_values

    def _handle_binary_op(self, op: pypcode.PcodeOp, op_symbol: str, signed=False):
        left = self.handle_get(op.inputs[0])
        right = self.handle_get(op.inputs[1])
        output = op.output

        if output is None:
            raise ValueError("Output of binary operation is None")

        self.handle_put(output, BinaryOp.create_binop(left, right, op_symbol, signed))

    def _init_handlers(self):
        self._handlers.update(
            {
                # Binary Arithmetic Operations
                pypcode.OpCode.INT_ADD: partial(self._handle_binary_op, op_symbol="+"),
                pypcode.OpCode.INT_MULT: partial(self._handle_binary_op, op_symbol="*"),
                # Binary Bitshits Operations
                pypcode.OpCode.INT_LEFT: partial(self._handle_binary_op, op_symbol="<<"),
                pypcode.OpCode.INT_RIGHT: partial(self._handle_binary_op, op_symbol=">>"),
                pypcode.OpCode.INT_AND: partial(self._handle_binary_op, op_symbol="&"),
                pypcode.OpCode.INT_SUB: partial(self._handle_binary_op, op_symbol="-"),
                pypcode.OpCode.INT_XOR: partial(self._handle_binary_op, op_symbol="^"),
                pypcode.OpCode.INT_OR: partial(self._handle_binary_op, op_symbol="|"),
                # Binary Conditional Operations
                pypcode.OpCode.INT_EQUAL: partial(self._handle_binary_op, op_symbol="=="),
                pypcode.OpCode.INT_SLESS: partial(self._handle_binary_op, op_symbol="<"),
                pypcode.OpCode.INT_LESS: partial(self._handle_binary_op, op_symbol="<"),
                pypcode.OpCode.INT_SLESSEQUAL: partial(self._handle_binary_op, op_symbol="<=", signed=True),
                pypcode.OpCode.INT_NOTEQUAL: partial(self._handle_binary_op, op_symbol="!="),
                # Variable / Memory Operations
                pypcode.OpCode.COPY: self._handle_copy,
                pypcode.OpCode.STORE: self._handle_store,
                pypcode.OpCode.LOAD: self._handle_load,
                pypcode.OpCode.IMARK: self._handle_imark,
                # CodeFlow Operations
                pypcode.OpCode.CALL: self._handle_call,
                pypcode.OpCode.CALLIND: self._handle_callind,
                pypcode.OpCode.CALLOTHER: self._do_nothing,  # TODO: Think if required, example is MIPS `rdhwr` in `sshd` `fileno` Function
                pypcode.OpCode.CBRANCH: self._handle_cbranch,
                pypcode.OpCode.BRANCH: self._handle_branch,
                pypcode.OpCode.BRANCHIND: self._handle_branchind,
                pypcode.OpCode.RETURN: self._do_nothing,
                # Unary Operations
                pypcode.OpCode.INT_2COMP: self._handle_int_2comp,
                pypcode.OpCode.BOOL_NEGATE: self._handle_bool_negate,
                pypcode.OpCode.INT_NEGATE: self._handle_int_negate,
                # Byte Operations
                pypcode.OpCode.INT_SEXT: self._handle_int_sext,
                pypcode.OpCode.INT_ZEXT: self._handle_int_zext,
                pypcode.OpCode.SUBPIECE: self._handle_subpiece,
            }
        )

        self.__put_handlers.update(
            {
                "register": self._handle_put_register,
                "unique": self._handle_put_unique,
            }
        )

        self.__get_handlers.update(
            {
                "register": self._handle_get_register,
                "const": self._handle_get_const,
                "unique": self._handle_get_unique,
                "ram": self._handle_get_const,
            }
        )

    def _analyze_block(self, blk: FunctionBlock):
        self.current_blk = blk
        current_address = blk.start

        while current_address < blk.end:
            for op in self.bin_func.opcodes[current_address].ops:
                self._handlers[op.opcode](op)
            current_address += self.bin_func.opcodes[current_address].bytes_size

    def __clear_after_callsite(self, instruction_state: InstructionState) -> None:
        if instruction_state.last_callsite is None:
            return

        for arg in instruction_state.last_callsite.args.values():
            if not isinstance(arg, BinaryOp):
                continue

            if not isinstance(arg.left, Register) or not isinstance(arg.right, int):
                continue

            if arg.left.offset == self.project.arch_regs.stackpointer:
                instruction_state.stack.pop(ctypes.c_int32(arg.right).value, None)

        for reg in list(instruction_state.regs.keys()):
            if reg not in self.project.arch_regs.unaffected:
                del instruction_state.regs[reg]

        instruction_state.regs[self.project.arch_regs.ret] = instruction_state.last_callsite
        instruction_state.last_callsite = None

    def __merge_dicts(self, condsite, iftrue, iffalse):
        """
        Merge two dict-like states using ConditionalExpression when values differ.
        If a key has the same value in both branches, keep that value.
        If a key exists in only one branch, keep that value.
        """
        merged = {}

        for k in iftrue.keys() | iffalse.keys():  # union of keys
            t = iftrue.get(k)
            f = iffalse.get(k)

            if k in iftrue and k in iffalse:
                merged[k] = t if t == f else ConditionalExpression(condsite, t, f)
            else:
                # appears in only one dict
                merged[k] = t if k in iftrue else f

        return merged

    def _get_block_and_state_from_mark(self, mark: int) -> tuple[FunctionBlock, InstructionState]:
        blk = self.bin_func.blocks_dict[mark]
        state = self.instructions_state[mark]
        return blk, state.goto_state.get(self.current_inst, state)

    def _handle_first_block(self):
        self.instructions_state[self.current_inst] = InstructionState()

    def _handle_one_previous_parent(self):
        previous_state = self.instructions_state[self.previous_marks[0]]
        current_state = previous_state.goto_state.get(self.current_inst, previous_state).copy()
        self.__clear_after_callsite(current_state)
        self.instructions_state[self.current_inst] = current_state

    def _merge_instruction_states(
        self, condsite: ConditionalSite, iftrue: InstructionState, iffalse: InstructionState
    ) -> InstructionState:
        """
        Helper to merge two states into a new state based on a conditional site.
        """
        merged_state = InstructionState()
        merged_state.regs = self.__merge_dicts(condsite, iftrue.regs, iffalse.regs)
        merged_state.unique = self.__merge_dicts(condsite, iftrue.unique, iffalse.unique)
        merged_state.ram = self.__merge_dicts(condsite, iftrue.ram, iffalse.ram)
        merged_state.stack = self.__merge_dicts(condsite, iftrue.stack, iffalse.stack)
        merged_state.used_arguments = set.union(iftrue.used_arguments, iffalse.used_arguments)
        return merged_state

    def _handle_imark(self, op: pypcode.PcodeOp):
        self.current_inst = op.inputs[0].offset
        self.__unfinished_condsite = None
        self.__conditional_move_condition = None

        if self.current_inst in self.loops_dict_start_address:
            good_marks = []
            loops = self.loops_dict[self.current_inst]
            loops_for_marks = loops.copy()

            if len(loops) == 2 and loops[0].start != loops[1].start:
                if loops[0].blocks & loops[1].blocks == loops[0].blocks:
                    loops_for_marks = [loops[0]]
                elif loops[0].blocks & loops[1].blocks == loops[1].blocks:
                    loops_for_marks = [loops[1]]

            for mark in self.previous_marks:
                # keep mark if it's NOT inside any of the loops starting at current_inst
                if not any(self.bin_func.blocks_dict[mark].start in loop.blocks for loop in loops_for_marks):
                    good_marks.append(mark)
            self.previous_marks = good_marks

        if len(self.previous_marks) == 0:  # first block
            self._handle_first_block()

        elif len(self.previous_marks) == 1:
            self._handle_one_previous_parent()

        else:
            # --- Generic N-Parent Merge Logic ---

            # 1. Collect all parent contexts (Block, State)
            # We copy the state to ensure we don't mutate the original instruction state
            # when clearing callsites.
            parent_contexts = []
            for mark in self.previous_marks:
                blk, state = self._get_block_and_state_from_mark(mark)
                state_copy = state.copy()
                self.__clear_after_callsite(state_copy)
                parent_contexts.append((blk, state_copy))

            # 2. Iteratively merge
            # We treat the first parent as the 'accumulator'
            current_blk, current_state = parent_contexts[0]

            for next_blk, next_state in parent_contexts[1:]:
                common_ancestor_addr = self.bin_func.common_ancestor(current_blk, next_blk)
                common_condsite = self.addr_to_codeflow_conditional_site.get(common_ancestor_addr, None)

                if common_condsite:
                    # STRICTLY MATCHING ORIGINAL LOGIC:
                    # Original: if self.bin_func.is_ancestor(blk_a.start, common_condsite.iffalse):
                    # We map blk_a -> current_blk (x), blk_b -> next_blk (y)

                    if self.bin_func.is_ancestor(current_blk.start, common_condsite.iffalse):
                        iftrue_state, iffalse_state = current_state, next_state
                    else:
                        iffalse_state, iftrue_state = current_state, next_state

                    current_state = self._merge_instruction_states(common_condsite, iftrue_state, iffalse_state)
                current_blk = next_blk

            self.instructions_state[self.current_inst] = current_state

        if self.current_inst in self.loops_dict_start_address:
            self._clear_looped_registers()

        self.previous_marks = [self.current_inst]

    def _clear_looped_registers(self):
        for loop in self.loops_dict[self.current_inst]:
            for blk in loop.blocks:
                current_blk = self.bin_func.blocks_dict[blk]
                current_address = current_blk.start
                while current_address < current_blk.end:
                    for op in self.bin_func.opcodes[current_address].ops:
                        if op.output is None or op.output.space.name != "register":
                            continue

                        self.instructions_state[self.current_inst].regs.pop(op.output.offset, None)
                        self.instructions_state[self.current_inst].used_arguments.add(op.output.offset)
                    current_address += self.bin_func.opcodes[current_address].bytes_size

    def _handle_copy(self, op: pypcode.PcodeOp):
        self.handle_put(op.output, self.handle_get(op.inputs[0]))

    def _handle_store(self, op: pypcode.PcodeOp):
        space = op.inputs[0].getSpaceFromConst().name
        offset = self.handle_get(op.inputs[1])
        val = self.handle_get(op.inputs[2])

        if space != "ram":
            raise NotImplementedError(f"STORE to space '{space}' is not implemented yet.")

        if isinstance(offset, int):
            self.instructions_state[self.current_inst].ram[offset] = val
            return

        if isinstance(offset, (CallSite, Arg)):
            self.memory_accesses.append(MemoryAccess(self.current_inst, offset, 0, MemoryAccessType.STORE, val))
            return

        if not isinstance(offset, BinaryOp):
            raise NotImplementedError(f"STORE with offset of type '{type(offset)}' is not implemented yet.")

        right = offset.right
        left = offset.left

        if not isinstance(left, Register) or left.offset != self.project.arch_regs.stackpointer:
            self.memory_accesses.append(MemoryAccess(self.current_inst, left, right, MemoryAccessType.STORE, val))
            return

        if isinstance(right, int) and offset.op == "+":
            signed_offset = ctypes.c_int32(right).value
            self.instructions_state[self.current_inst].stack[signed_offset] = val

    def _handle_int_2comp(self, op: pypcode.PcodeOp):
        val = op.inputs[0].offset  # INT_2COMP is always int values.

        self.handle_put(op.output, -val)

    def _handle_call(self, op: pypcode.PcodeOp):
        self._handle_callind(op)

    def _create_callsite(self, target: Any) -> CallSite:
        resolved_target = self.__extract_target(target)
        args = self.__collect_call_arguments()

        callsite = CallSite(self.current_inst, resolved_target, frozendict(args))
        self.__register_callsite(callsite)

        return callsite

    def _handle_callind(self, op: pypcode.PcodeOp):
        target = self.handle_get(op.inputs[0])

        self._create_callsite(target)

    def _handle_int_zext(self, op: pypcode.PcodeOp):
        # Assuming zext(X) == X for now.
        self.handle_put(op.output, self.handle_get(op.inputs[0]))

    def _handle_subpiece(self, op: pypcode.PcodeOp):
        # Assumming X = SUBPIECE(X, N) for now
        self.handle_put(op.output, self.handle_get(op.inputs[0]))

    def _create_condsite(self, condition: BinaryOp, goto_iftrue: int, goto_iffalse: int) -> ConditionalSite:
        condsite = ConditionalSite(self.current_inst, condition, goto_iftrue, goto_iffalse)
        self.conditional_sites.append(condsite)
        if goto_iftrue != 2:
            self.addr_to_codeflow_conditional_site[self.current_blk.start] = condsite

        loops = self.loops_dict.get(self.current_inst, None)
        if loops is not None and isinstance(condition, BinaryOp):
            for loop in loops:
                if goto_iftrue not in loop.blocks:
                    loop.exit_conditions[self.current_inst] = condition
                elif goto_iffalse not in loop.blocks:
                    loop.exit_conditions[self.current_inst] = condition.negate()

        return condsite

    def _handle_branch(self, op: pypcode.PcodeOp):
        goto_addr = self.handle_get(op.inputs[0])

        if goto_addr < self.bin_func.start or goto_addr >= self.bin_func.end:
            self._create_callsite(goto_addr)
            return

        if self.__unfinished_condsite is None:
            return

        self._create_condsite(self.__unfinished_condsite.condition, self.__unfinished_condsite.goto_iftrue, goto_addr)

    def _handle_cbranch(self, op: pypcode.PcodeOp):
        goto_iftrue = self.handle_get(op.inputs[0])
        goto_iffalse = self.bin_func._next_address(self.current_inst)
        condition: BinaryOp | int = self.handle_get(op.inputs[1])

        if isinstance(condition, int):  # Sometimes the result is known at analysis time, for example LL/SC instructions
            return

        if goto_iftrue == 2:
            self.__conditional_move_condition = self._create_condsite(condition, goto_iftrue, goto_iffalse)
            return

        if goto_iffalse == goto_iftrue:
            self.instructions_state[self.current_inst].goto_state[goto_iftrue] = self.instructions_state[
                self.current_inst
            ].copy()
            self.__unfinished_condsite = _UnfinishedConditionalSite(condition, goto_iftrue)
            return

        self._create_condsite(condition, goto_iftrue, goto_iffalse)

    def _do_nothing(self, op: pypcode.PcodeOp):
        pass

    def _handle_branchind(self, op: pypcode.PcodeOp):
        target = self.handle_get(op.inputs[0])

        self._create_callsite(target)

    def _handle_bool_negate(self, op: pypcode.PcodeOp):
        bool_expr: BinaryOp | int = self.handle_get(op.inputs[0])

        if isinstance(bool_expr, int):
            self._handle_int_negate(op)
            return

        self.handle_put(op.output, bool_expr.negate())

    def _handle_int_negate(self, op: pypcode.PcodeOp):
        int_expr: BinaryOp = self.handle_get(op.inputs[0])

        self.handle_put(op.output, UnaryOp(int_expr, "~"))

    def _handle_int_sext(self, op: pypcode.PcodeOp):
        # TODO: For now assuming X = sext(X)
        self.handle_put(op.output, self.handle_get(op.inputs[0]))

    def _handle_load(self, op: pypcode.PcodeOp):
        arch = self.project.arch_regs
        state = self.instructions_state[self.current_inst]

        space = op.inputs[0].getSpaceFromConst().name
        offset = self.handle_get(op.inputs[1])

        if space != "ram":
            raise NotImplementedError(f"LOAD to space '{space}' is not implemented yet.")

        if isinstance(offset, (int, UnaryOp)):
            self.handle_put(op.output, UnaryOp(offset, "*"))
            return

        if not isinstance(offset, BinaryOp):
            res = MemoryAccess(self.current_inst, offset, 0, MemoryAccessType.LOAD)
            self.handle_put(op.output, res)
            self.memory_accesses.append(res)
            return

        left = offset.left
        right = offset.right

        if not isinstance(left, Register) or left.offset != self.project.arch_regs.stackpointer:
            res = MemoryAccess(self.current_inst, left, right, MemoryAccessType.LOAD)
            self.handle_put(op.output, res)
            self.memory_accesses.append(res)
            return

        if isinstance(left, Register) and left.offset == arch.stackpointer:
            signed_value = ctypes.c_int32(right).value
            if signed_value >= arch.stack_argument_offset:
                res = Arg((signed_value) // 4)
            else:
                res = state.stack.get(signed_value, None)
                if res is None:
                    unsigned_value = ctypes.c_uint32(right).value
                    res = MemoryAccess(
                        self.current_inst, self._first_stack_access, unsigned_value, MemoryAccessType.LOAD
                    )
                    state.stack[signed_value] = res

            self.handle_put(op.output, res)

    def _handle_get_register(self, offset: int) -> Any:
        state = self.instructions_state[self.current_inst]
        arch = self.project.arch_regs

        reg = state.regs.get(offset, None)
        if reg is not None:
            return reg

        res: Register | Arg
        if offset in arch.rev_arguments and state.last_callsite is None and offset not in state.used_arguments:
            res = Arg(arch.rev_arguments[offset])
        else:
            res = Register(offset, self.current_inst, self.project)
            if offset == arch.stackpointer and self._first_stack_access is None:
                self._first_stack_access = res

        state.regs[offset] = res
        return res

    def _handle_get_const(self, offset: int) -> Any:
        return offset

    def _handle_get_unique(self, offset: int) -> Any:
        return self.instructions_state[self.current_inst].unique.get(offset, None)

    def handle_get(self, input: pypcode.Varnode | FakeVarnode) -> Any:
        return self.__get_handlers[input.space.name](input.offset)

    def _handle_put_register(self, offset: int, val: Any):
        if self.__conditional_move_condition is None:
            self.instructions_state[self.current_inst].regs[offset] = val
        else:
            self.instructions_state[self.current_inst].regs[offset] = ConditionalExpression(
                self.__conditional_move_condition,
                self.instructions_state[self.current_inst].regs[offset],
                val,
            )

    def _handle_put_unique(self, offset: int, val: Any):
        self.instructions_state[self.current_inst].unique[offset] = val

    def handle_put(self, output: pypcode.Varnode | None, val: Any):
        self.__put_handlers[output.space.name](output.offset, val)

    def __extract_target(self, target: Any) -> Any:
        if self.project.arch_regs.does_isa_switches and isinstance(target, BinaryOp):
            target = target.right

        return target

    def __register_callsite(self, callsite: CallSite) -> None:
        self.callsites.append(callsite)
        self.instructions_state[self.current_inst].last_callsite = callsite

    def __collect_call_arguments(self) -> dict[int, Any]:
        state = self.instructions_state[self.current_inst]
        arch = self.project.arch_regs

        args = {arg_num: state.regs[reg] for arg_num, reg in arch.arguments.items() if reg in state.regs}

        current_stack = state.regs.get(arch.stackpointer, None)
        if current_stack is not None and (len(args) == len(arch.arguments) or len(args) == 0):
            if not isinstance(current_stack, BinaryOp):
                raise ValueError("Current stack pointer is expected to be a BinaryOp")

            stack_argument_offset = ctypes.c_int32(current_stack.right + arch.stack_argument_offset).value
            arg_num = len(arch.arguments)

            while stack_argument_offset in state.stack:
                args[arg_num] = state.stack[stack_argument_offset]
                stack_argument_offset += 0x4
                arg_num += 1

        max_reg_arg = min(max(args.keys(), default=0), len(arch.arguments))

        for arg_num, _ in arch.arguments.items():
            if arg_num not in args and arg_num < max_reg_arg:
                args[arg_num] = self.handle_get(FakeVarnode(FakeAddrSpace("register"), arch.arguments[arg_num], 4))

        return args


class FakeAddrSpace:
    def __init__(self, name: str):
        self.name = name


class FakeVarnode:
    def __init__(self, space: FakeAddrSpace, offset: int, size: int):
        self.space = space
        self.offset = offset
        self.size = size


@dataclass
class _UnfinishedConditionalSite:
    condition: BinaryOp
    goto_iftrue: int