angr 9.2.192__cp311-cp311-macosx_10_12_x86_64.whl

angr/analyses/decompiler/optimization_passes/win_stack_canary_simplifier.py

@@ -0,0 +1,477 @@
+# pylint:disable=too-many-boolean-expressions
+from __future__ import annotations
+from collections import defaultdict
+import logging
+
+import angr.ailment as ailment
+import cle
+
+from angr.utils.funcid import is_function_security_check_cookie
+from angr.analyses.decompiler.stack_item import StackItem, StackItemType
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+
+_l = logging.getLogger(name=__name__)
+
+
+class WinStackCanarySimplifier(OptimizationPass):
+    """
+    Removes stack canary checks from decompilation results for Windows PE files.
+
+    we need to run this pass before performing any full-function simplification. Otherwise the effects of
+    _security_cookie will be propagated.
+    """
+
+    ARCHES = [
+        "X86",
+        "AMD64",
+    ]
+    PLATFORMS = ["windows"]
+    STAGE = OptimizationPassStage.AFTER_SINGLE_BLOCK_SIMPLIFICATION
+    NAME = "Simplify stack canaries in Windows PE files"
+    DESCRIPTION = __doc__.strip()  # type:ignore
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+        self._security_cookie_addr: int | None = None
+        if isinstance(self.project.loader.main_object, cle.PE):
+            self._security_cookie_addr = self.project.loader.main_object.load_config.get("SecurityCookie", None)
+            if self._security_cookie_addr is None:
+                # maybe it's just not set - check labels
+                self._security_cookie_addr = self.project.kb.labels.lookup("_security_cookie")
+
+        self.analyze()
+
+    def _check(self):
+        if self._security_cookie_addr is None:
+            return False, None
+
+        # Check the first block and see if there is any statement reading data from _security_cookie
+        init_stmts = self._find_canary_init_stmt()
+
+        return init_stmts is not None, {"init_stmts": init_stmts}
+
+    def _analyze(self, cache=None):
+        init_stmts = None
+        if cache is not None:
+            init_stmts = cache.get("init_stmts", None)
+
+        if init_stmts is None:
+            init_stmts = self._find_canary_init_stmt()
+
+        if init_stmts is None:
+            return
+
+        # Look for the statement that loads back canary value from the stack
+        first_block, canary_init_stmt_ids = init_stmts
+        canary_init_stmt = first_block.statements[canary_init_stmt_ids[-1]]
+        # where is the stack canary stored?
+        if not isinstance(canary_init_stmt, ailment.Stmt.Store):
+            return
+        store_offset = self._get_bp_offset(canary_init_stmt.addr, canary_init_stmt.tags["ins_addr"])
+        if store_offset is None:
+            _l.debug(
+                "Unsupported canary storing location %s. Expects a StackBaseOffset or (bp - Const).",
+                canary_init_stmt.addr,
+            )
+            return
+
+        # The function should have at least one end point calling _security_check_cookie
+        # note that (at least for now) we rely on FLIRT to identify the _security_check_cookie function inside the
+        # binary.
+        # TODO: Add function matching logic to this simplifier
+
+        # Find all nodes with 0 out-degrees
+        all_endpoint_addrs = [node.addr for node in self._func.graph.nodes() if self._func.graph.out_degree(node) == 0]
+
+        # Before node duplication, each pair of canary-check-success and canary-check-failure nodes have a common
+        # predecessor.
+        # map endpoint addrs to their common predecessors
+        pred_addr_to_endpoint_addrs: dict[int, set[int]] = defaultdict(set)
+        for node_addr in all_endpoint_addrs:
+            preds = self._func.graph.predecessors(self._func.get_node(node_addr))
+            for pred in preds:
+                pred_addr_to_endpoint_addrs[pred.addr].add(node_addr)
+
+        found_endpoints = False
+        for pred_addr in pred_addr_to_endpoint_addrs:
+            # the predecessor should call _security_check_cookie
+            endpoint_preds = list(self._get_blocks(pred_addr))
+            if not endpoint_preds:
+                continue
+            if self._find_stmt_calling_security_check_cookie(endpoint_preds[0]) is None:
+                _l.debug("The predecessor does not call _security_check_cookie().")
+                continue
+
+            nodes_to_process = []
+            for pred in endpoint_preds:
+                check_call_stmt_idx = self._find_stmt_calling_security_check_cookie(pred)
+                if check_call_stmt_idx is None:
+                    _l.debug("Cannot find the statement calling _security_check_cookie() in the predecessor.")
+                    continue
+
+                canary_storing_stmt_idx = (
+                    self._find_amd64_canary_storing_stmt(pred, store_offset)
+                    if self.project.arch.name == "AMD64"
+                    else self._find_x86_canary_storing_stmt(pred, store_offset)
+                )
+                if canary_storing_stmt_idx is None:
+                    _l.debug("Cannot find the canary check statement in the predecessor.")
+                    continue
+
+                return_addr_storing_stmt_idx = self._find_return_addr_storing_stmt(pred)
+                if return_addr_storing_stmt_idx is None:
+                    _l.debug("Cannot find the return address storing statement in the predecessor.")
+                    continue
+
+                nodes_to_process.append(
+                    (pred, check_call_stmt_idx, canary_storing_stmt_idx, return_addr_storing_stmt_idx)
+                )
+
+            # Now patch this function.
+            for pred, check_call_stmt_idx, canary_storing_stmt_idx, return_addr_storing_stmt_idx in nodes_to_process:
+                # Patch the pred so that it jumps to the one that is not stack_chk_fail_caller
+                stmts = []
+                for stmt_idx, stmt in enumerate(pred.statements):
+                    if stmt_idx in {check_call_stmt_idx, canary_storing_stmt_idx, return_addr_storing_stmt_idx}:
+                        continue
+                    stmts.append(stmt)
+                pred_copy = pred.copy(statements=stmts)
+
+                self._update_block(pred, pred_copy)
+
+                found_endpoints = True
+
+        if found_endpoints:
+            # Remove the statement that loads the stack canary from fs
+            first_block_copy = first_block.copy()
+            for stmt_idx in sorted(canary_init_stmt_ids, reverse=True):
+                first_block_copy.statements.pop(stmt_idx)
+            self._update_block(first_block, first_block_copy)
+
+            # update stack_items
+            self.stack_items[store_offset] = StackItem(
+                store_offset, canary_init_stmt.size, "canary", StackItemType.STACK_CANARY
+            )
+
+    def _find_canary_init_stmt(self) -> tuple[ailment.Block, list[int]] | None:
+        first_block = self._get_block(self._func.addr)
+        if first_block is None:
+            return None
+
+        r = self._extract_canary_init_stmt_from_block(first_block)
+        if r is not None:
+            return first_block, r
+
+        # if the first block is calling alloca_probe, we may want to take the second block instead
+        if (
+            first_block.statements
+            and first_block.original_size > 0
+            and isinstance(first_block.statements[-1], ailment.statement.Call)
+            and isinstance(first_block.statements[-1].target, ailment.expression.Const)
+        ):
+            # check if the target is alloca_probe
+            callee_addr = first_block.statements[-1].target.value
+            if (
+                self.kb.functions.contains_addr(callee_addr)
+                and self.kb.functions.get_by_addr(callee_addr).info.get("is_alloca_probe", False) is True
+            ):
+                second_block = self._get_block(first_block.addr + first_block.original_size)
+                if second_block is not None:
+                    r = self._extract_canary_init_stmt_from_block(second_block)
+                    if r is not None:
+                        return second_block, r
+        return None
+
+    def _extract_canary_init_stmt_from_block(self, block: ailment.Block) -> list[int] | None:
+        load_stmt_idx = None
+        load_reg = None
+        xor_stmt_idx = None
+        xored_reg = None
+
+        assert self._security_cookie_addr is not None
+
+        for idx, stmt in enumerate(block.statements):
+            # if we are lucky and things get folded into one statement:
+            if (
+                isinstance(stmt, ailment.Stmt.Store)
+                and isinstance(stmt.addr, ailment.Expr.StackBaseOffset)
+                and self._is_expr_loading_stack_cookie(stmt.data, self._security_cookie_addr)
+            ):
+                return [idx]
+            # or if we are unlucky and the load and the xor are two different statements
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.was_reg
+                and isinstance(stmt.src, ailment.Expr.Load)
+                and isinstance(stmt.src.addr, ailment.Expr.Const)
+            ):
+                load_addr = stmt.src.addr.value
+                if load_addr == self._security_cookie_addr:
+                    load_stmt_idx = idx
+                    load_reg = stmt.dst.reg_offset
+            if load_stmt_idx is not None and load_reg is not None and xor_stmt_idx is None and idx >= load_stmt_idx + 1:
+                assert self.project.arch.bp_offset is not None
+                if (
+                    isinstance(stmt, ailment.Stmt.Assignment)
+                    and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                    and stmt.dst.was_reg
+                    and not self.project.arch.is_artificial_register(stmt.dst.reg_offset, stmt.dst.size)
+                    and self._is_expr_xoring_stack_cookie_reg(stmt.src, load_reg, self.project.arch.bp_offset)
+                ):
+                    xor_stmt_idx = idx
+                    xored_reg = stmt.dst.reg_offset
+            if load_stmt_idx is not None and xor_stmt_idx is not None and idx == xor_stmt_idx + 1:
+                if (
+                    isinstance(stmt, ailment.Stmt.Store)
+                    and (
+                        isinstance(stmt.addr, ailment.Expr.StackBaseOffset)
+                        or (
+                            isinstance(stmt.addr, ailment.Expr.BinaryOp)
+                            and stmt.addr.op == "Sub"
+                            and isinstance(stmt.addr.operands[0], ailment.Expr.VirtualVariable)
+                            and stmt.addr.operands[0].was_reg
+                            and stmt.addr.operands[0].reg_offset == self.project.arch.registers["ebp"][0]
+                            and isinstance(stmt.addr.operands[1], ailment.Expr.Const)
+                        )
+                    )
+                    and isinstance(stmt.data, ailment.Expr.VirtualVariable)
+                    and stmt.data.was_reg
+                    and stmt.data.reg_offset == xored_reg
+                ):
+                    return [load_stmt_idx, xor_stmt_idx, idx]
+                break
+            if load_stmt_idx is not None and xor_stmt_idx is None and idx >= load_stmt_idx + 1:  # noqa:SIM102
+                if isinstance(stmt, ailment.Stmt.Store) and (
+                    isinstance(stmt.addr, ailment.Expr.StackBaseOffset)
+                    or (
+                        isinstance(stmt.addr, ailment.Expr.BinaryOp)
+                        and stmt.addr.op == "Sub"
+                        and isinstance(stmt.addr.operands[0], ailment.Expr.VirtualVariable)
+                        and stmt.addr.operands[0].was_reg
+                        and stmt.addr.operands[0].reg_offset == self.project.arch.registers["ebp"][0]
+                        and isinstance(stmt.addr.operands[1], ailment.Expr.Const)
+                    )
+                ):
+                    if (
+                        isinstance(stmt.data, ailment.Expr.VirtualVariable)
+                        and stmt.data.was_reg
+                        and stmt.data.reg_offset == load_reg
+                    ):
+                        return [load_stmt_idx, idx]
+                    break
+        return None
+
+    def _find_amd64_canary_storing_stmt(self, block, canary_value_stack_offset):
+        load_stmt_idx = None
+
+        for idx, stmt in enumerate(block.statements):
+            # when we are lucky, we have one instruction
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.was_reg
+                and stmt.dst.reg_offset == self.project.arch.registers["rcx"][0]
+            ):
+                if isinstance(stmt.src, ailment.Expr.BinaryOp) and stmt.src.op == "Xor":
+                    op0, op1 = stmt.src.operands
+                    if (
+                        isinstance(op0, ailment.Expr.Load)
+                        and isinstance(op0.addr, ailment.Expr.StackBaseOffset)
+                        and op0.addr.offset == canary_value_stack_offset
+                    ) and isinstance(op1, ailment.Expr.StackBaseOffset):
+                        # found it
+                        return idx
+                elif isinstance(stmt.src, ailment.Expr.Load):
+                    if (
+                        isinstance(stmt.src.addr, ailment.Expr.StackBaseOffset)
+                        and stmt.src.addr.offset == canary_value_stack_offset
+                    ):
+                        # found it
+                        return idx
+            # or when we are unlucky, we have two instructions...
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.reg_offset == self.project.arch.registers["rcx"][0]
+                and isinstance(stmt.src, ailment.Expr.Load)
+                and isinstance(stmt.src.addr, ailment.Expr.StackBaseOffset)
+                and stmt.src.addr.offset == canary_value_stack_offset
+            ):
+                load_stmt_idx = idx
+            if (
+                load_stmt_idx is not None
+                and idx >= load_stmt_idx + 1
+                and (
+                    isinstance(stmt, ailment.Stmt.Assignment)
+                    and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                    and stmt.dst.was_reg
+                    and isinstance(stmt.src, ailment.Expr.BinaryOp)
+                    and stmt.src.op == "Xor"
+                )
+                and (
+                    isinstance(stmt.src.operands[0], ailment.Expr.VirtualVariable)
+                    and stmt.src.operands[0].was_reg
+                    and stmt.src.operands[0].reg_offset == self.project.arch.registers["rcx"][0]
+                    and isinstance(stmt.src.operands[1], ailment.Expr.StackBaseOffset)
+                )
+            ):
+                return idx
+        return None
+
+    def _find_x86_canary_storing_stmt(self, block, canary_value_stack_offset):
+        load_stmt_idx = None
+        canary_reg_dst_offset = None
+
+        for idx, stmt in enumerate(block.statements):
+            # when we are lucky, we have one instruction
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.was_reg
+                and not self.project.arch.is_artificial_register(stmt.dst.reg_offset, stmt.dst.size)
+            ):
+                if isinstance(stmt.src, ailment.Expr.BinaryOp) and stmt.src.op == "Xor":
+                    op0, op1 = stmt.src.operands
+                    if (
+                        isinstance(op0, ailment.Expr.Load)
+                        and self._get_bp_offset(op0.addr, stmt.tags["ins_addr"]) == canary_value_stack_offset
+                    ) and isinstance(op1, ailment.Expr.StackBaseOffset):
+                        # found it
+                        return idx
+                elif isinstance(stmt.src, ailment.Expr.Load):
+                    if (
+                        isinstance(stmt.src.addr, ailment.Expr.StackBaseOffset)
+                        and stmt.src.addr.offset == canary_value_stack_offset
+                    ):
+                        # found it
+                        return idx
+            # or when we are unlucky, we have two instructions...
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and not self.project.arch.is_artificial_register(stmt.dst.reg_offset, stmt.dst.size)
+                and isinstance(stmt.src, ailment.Expr.Load)
+                and self._get_bp_offset(stmt.src.addr, stmt.tags["ins_addr"]) == canary_value_stack_offset
+            ):
+                load_stmt_idx = idx
+                canary_reg_dst_offset = stmt.dst.reg_offset
+            if (
+                load_stmt_idx is not None
+                and idx >= load_stmt_idx + 1
+                and (
+                    isinstance(stmt, ailment.Stmt.Assignment)
+                    and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                    and stmt.dst.was_reg
+                    and not self.project.arch.is_artificial_register(stmt.dst.reg_offset, stmt.dst.size)
+                    and isinstance(stmt.src, ailment.Expr.BinaryOp)
+                    and stmt.src.op == "Xor"
+                )
+                and (
+                    isinstance(stmt.src.operands[0], ailment.Expr.VirtualVariable)
+                    and stmt.src.operands[0].was_reg
+                    and stmt.src.operands[0].reg_offset == canary_reg_dst_offset
+                    and self._get_bp_offset(stmt.src.operands[1], stmt.tags["ins_addr"]) is not None
+                )
+            ):
+                return idx
+        return None
+
+    def _get_bp_offset(self, expr: ailment.Expr.Expression, ins_addr: int) -> int | None:
+        if isinstance(expr, ailment.Expr.StackBaseOffset):
+            return expr.offset
+        if (
+            isinstance(expr, ailment.Expr.VirtualVariable)
+            and expr.was_reg
+            and expr.reg_offset == self.project.arch.bp_offset
+        ):
+            if self._stack_pointer_tracker is None:
+                return None
+            return self._stack_pointer_tracker.offset_before(ins_addr, self.project.arch.bp_offset)
+        if (
+            isinstance(expr, ailment.Expr.BinaryOp)
+            and expr.op == "Sub"
+            and isinstance(expr.operands[0], ailment.Expr.VirtualVariable)
+            and expr.operands[0].was_reg
+            and expr.operands[0].reg_offset == self.project.arch.bp_offset
+            and isinstance(expr.operands[1], ailment.Expr.Const)
+        ):
+            if self._stack_pointer_tracker is None:
+                return None
+            base_bp_offset = self._stack_pointer_tracker.offset_before(ins_addr, self.project.arch.bp_offset)
+            if base_bp_offset is None:
+                return None
+            bp_off = base_bp_offset - expr.operands[1].value
+            mask = 0xFFFF_FFFF if self.project.arch.bits == 32 else 0xFFFF_FFFF_FFFF_FFFF
+            return bp_off & mask
+        return None
+
+    @staticmethod
+    def _find_return_addr_storing_stmt(block):
+        for idx, stmt in enumerate(block.statements):
+            if (
+                isinstance(stmt, ailment.Stmt.Store)
+                and isinstance(stmt.addr, ailment.Expr.StackBaseOffset)
+                and isinstance(stmt.data, ailment.Expr.Const)
+                and stmt.data.value == block.addr + block.original_size
+            ):
+                return idx
+        return None
+
+    def _find_stmt_calling_security_check_cookie(self, node):
+        assert self._security_cookie_addr is not None
+        for idx, stmt in enumerate(node.statements):
+            if isinstance(stmt, ailment.Stmt.Call) and isinstance(stmt.target, ailment.Expr.Const):
+                const_target = stmt.target.value
+                if self.kb.functions.contains_addr(const_target):
+                    func = self.kb.functions.get_by_addr(const_target)
+                    assert func is not None
+                    if func.name == "_security_check_cookie" or is_function_security_check_cookie(
+                        func, self.project, self._security_cookie_addr
+                    ):
+                        return idx
+
+        return None
+
+    @staticmethod
+    def _is_expr_loading_stack_cookie(expr: ailment.Expr.Expression, security_cookie_addr: int) -> bool:
+        if (
+            isinstance(expr, ailment.Expr.BinaryOp)
+            and expr.op == "Xor"
+            and isinstance(expr.operands[1], ailment.Expr.StackBaseOffset)
+            and isinstance(expr.operands[0], ailment.Expr.Load)
+            and isinstance(expr.operands[0].addr, ailment.Expr.Const)
+        ):
+            # Check addr: must be __security_cookie
+            load_addr = expr.operands[0].addr.value
+            if load_addr == security_cookie_addr:
+                return True
+
+        if isinstance(expr, ailment.Expr.Load) and isinstance(expr.addr, ailment.Expr.Const):
+            load_addr = expr.addr.value
+            # Check addr: must be __security_cookie
+            if load_addr == security_cookie_addr:
+                return True
+
+        return False
+
+    @staticmethod
+    def _is_expr_xoring_stack_cookie_reg(
+        expr: ailment.Expr.Expression, security_cookie_reg: int, bp_offset: int
+    ) -> bool:
+        return (
+            isinstance(expr, ailment.Expr.BinaryOp)
+            and expr.op == "Xor"
+            and isinstance(expr.operands[0], ailment.Expr.VirtualVariable)
+            and expr.operands[0].was_reg
+            and expr.operands[0].reg_offset == security_cookie_reg
+            and (
+                isinstance(expr.operands[1], ailment.Expr.StackBaseOffset)
+                or (
+                    isinstance(expr.operands[1], ailment.Expr.VirtualVariable)
+                    and expr.operands[1].was_reg
+                    and expr.operands[1].reg_offset == bp_offset
+                )
+            )
+        )

angr/analyses/decompiler/optimization_passes/x86_gcc_getpc_simplifier.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+from typing import Any
+import logging
+
+import angr.ailment as ailment
+
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+_l = logging.getLogger(name=__name__)
+
+
+class X86GccGetPcSimplifier(OptimizationPass):
+    """
+    Simplifies __x86.get_pc_thunk calls.
+    """
+
+    ARCHES = ["X86"]
+    PLATFORMS = ["linux"]
+    STAGE = OptimizationPassStage.BEFORE_SSA_LEVEL0_TRANSFORMATION
+    NAME = "Simplify getpc()"
+    DESCRIPTION = __doc__.strip()
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+        self.analyze()
+
+    def _check(self):
+        getpc_calls = self._find_getpc_calls()
+
+        return bool(getpc_calls), {
+            "getpc_calls": getpc_calls,
+        }
+
+    def _analyze(self, cache=None):
+        getpc_calls = None
+
+        if cache is not None:
+            getpc_calls = cache.get("getpc_calls", None)
+
+        if getpc_calls is None:
+            getpc_calls = self._find_getpc_calls()
+
+        if not getpc_calls:
+            return
+
+        # update each block
+        for block_key, stmt_idx, getpc_reg, getpc_reg_value in getpc_calls:
+            pcreg_offset = self.project.arch.registers[getpc_reg][0]
+
+            old_block = self.blocks_by_addr_and_idx[block_key]
+            block = old_block.copy()
+            old_stmt = block.statements[stmt_idx]
+            block.statements[stmt_idx] = ailment.Stmt.Assignment(
+                old_stmt.idx,
+                ailment.Expr.Register(None, None, pcreg_offset, 32, reg_name=getpc_reg),
+                ailment.Expr.Const(None, None, getpc_reg_value, 32),
+                **old_stmt.tags,
+            )
+            # remove the statement that pushes return address onto the stack
+            if stmt_idx > 0 and isinstance(block.statements[stmt_idx - 1], ailment.Stmt.Store):
+                block.statements = block.statements[: stmt_idx - 1] + block.statements[stmt_idx:]
+            self._update_block(old_block, block)
+
+    def _find_getpc_calls(self) -> list[tuple[Any, int, str, int]]:
+        """
+        Find all blocks that are calling __x86.get_pc_thunk functions.
+
+        :return:    A list of tuples. Each tuple is in the form of (block_key, statement ID, pc-storing register,
+                    value of the pc-storing register).
+        """
+
+        results = []
+        for key, block in self._blocks_by_addr_and_idx.items():
+            if (
+                block.statements
+                and isinstance(block.statements[-1], ailment.Stmt.Call)
+                and isinstance(block.statements[-1].target, ailment.Expr.Const)
+            ):
+                call_func_addr = block.statements[-1].target.value
+                try:
+                    call_func = self.kb.functions.get_by_addr(call_func_addr)
+                except KeyError:
+                    continue
+                if "get_pc" in call_func.info:
+                    results.append(
+                        (key, len(block.statements) - 1, call_func.info["get_pc"], block.addr + block.original_size),
+                    )
+        return results