pcode-graph 1.0.0__py3-none-any.whl

pcode_graph/simplifier.py

@@ -0,0 +1,185 @@
+from pcode_graph.checker import check_graph
+from pcode_graph.graph import CDG, Edge, EdgeKinds, NodeIndex, OpCodes, NodeKinds
+
+from loguru import logger
+from pathlib import Path
+
+
+class GraphSimplifier:
+
+    def __init__(
+        self,
+        graph: CDG,
+    ) -> None:
+        self.graph = graph
+
+    def _update_proxy(self):
+        self._proxy = self.graph.compute_edge_proxy()
+
+    def _remove_node(self, index: NodeIndex):
+        def remap_index(n: NodeIndex):
+            assert n != index
+            return n - int(n > index)
+
+        # logger.debug(f"Remove {self.graph.nodes[index]} at {index}")
+
+        # Remap edges
+        proxy = self._proxy[index]
+
+        control_succs = proxy.get_control_successors()
+        assert (
+            len(control_succs) <= 1
+        ), "More than one control successor to the node to remove"
+
+        data_preds = proxy.get_data_predecessors()
+
+        # Convert to set to handle particular case of Phi node
+        assert len(data_preds) <= 1 or (
+            self.graph.nodes[index].kind == NodeKinds.Phi and len(set(data_preds)) == 1
+        ), "More than one data predecessor to the node to remove"
+
+        kept_edges: list[Edge] = []
+
+        def remap_edge(
+            kind: EdgeKinds, source: NodeIndex, dest: NodeIndex, operand: int | None
+        ):
+            # logger.debug(
+            #     f"Remap {self.graph.dump_edge(Edge(kind=kind, source_node=source, destination_node=dest, operand_number=operand))}"
+            # )
+            edge = Edge(
+                source_node=remap_index(source),
+                destination_node=remap_index(dest),
+                kind=kind,
+                operand_number=operand,
+            )
+            kept_edges.append(edge)
+
+        for e in self.graph.edges:
+            if e.source_node == index:
+                if e.kind == EdgeKinds.Data:
+                    remap_edge(
+                        e.kind,
+                        data_preds[0],
+                        e.destination_node,
+                        e.operand_number,
+                    )
+            elif e.destination_node == index:
+                if e.kind == EdgeKinds.Control:
+                    remap_edge(
+                        e.kind,
+                        e.source_node,
+                        control_succs[0],
+                        e.operand_number,
+                    )
+            else:
+                remap_edge(e.kind, e.source_node, e.destination_node, e.operand_number)
+
+        # Patch the graph
+        self.graph.edges = kept_edges
+        self.graph.nodes.pop(index)
+
+        self._graph_changed = True
+        self._update_proxy()
+
+        # self.debug_helper()
+
+    def _remove_edge(self, edge: Edge):
+
+        self.graph.remove_edge(edge)
+        self._graph_changed = True
+        self._update_proxy()
+
+    def debug_helper(self):
+        try:
+            self.rounds += 1
+        except AttributeError:
+            self.rounds = 0
+        path = Path(f"round{self.rounds}.md")
+        path.write_text(str(self.graph))
+        logger.debug(f"Written {path}")
+        check_graph(self.graph)
+
+    def process(self):
+
+        self._graph_changed = True
+        self._update_proxy()
+
+        while self._graph_changed:
+            self._graph_changed = False
+
+            # Remove useless phi nodes
+            for index, node in enumerate(self.graph.nodes):
+
+                if node.kind == NodeKinds.Phi:
+
+                    phi_data_preds = set(self._proxy[index].get_data_predecessors())
+                    if len(phi_data_preds) > 1:
+                        # Keep this Phi node
+                        continue
+
+                    self._remove_node(index)
+                    break
+
+                elif node.opcode == OpCodes.COPY:
+
+                    ctrl_preds = set(self._proxy[index].get_control_predecessors())
+                    ctrl_succs = self._proxy[index].get_control_successors()
+
+                    remove = True
+                    if ctrl_preds and ctrl_succs:
+                        (ctrl_succ,) = ctrl_succs
+
+                        remove = False
+                        # Remove copies that dominates their successor...
+                        if self._proxy[ctrl_succ].get_control_predecessors() == [index]:
+                            remove = True
+                        elif len(ctrl_preds) == 1:
+                            # ...or are dominated by their unique predecessor
+                            (ctrl_pred,) = ctrl_preds
+                            if self._proxy[ctrl_pred].get_control_successors() == [
+                                index
+                            ]:
+                                remove = True
+
+                    if remove:
+                        self._remove_node(index)
+                        break
+
+                elif node.kind == NodeKinds.Constant:                    
+                    if not self._proxy[index].get_data_successors():
+                        # Remove constants let orphan by other simplifications
+                        self._remove_node(index)
+                        break
+
+                elif node.opcode in {
+                    OpCodes.INT_ZEXT,
+                    OpCodes.INT_SEXT,
+                }:
+                    self._remove_node(index)
+                    break
+
+                elif node.opcode == OpCodes.SUBPIECE:
+
+                    operands: list[Edge] = self._proxy[index].get_inputs_edges()
+
+                    assert len(operands) == 2
+                    num_trunc_bytes = self.graph.nodes[operands[1].source_node]
+                    assert num_trunc_bytes.kind == NodeKinds.Constant
+                    if num_trunc_bytes.value == 0:
+                        self._remove_edge(operands[1])
+                        self._remove_node(index)                        
+                        break
+
+                elif node.opcode == OpCodes.INT_MULT:
+                    # Remove mulitiplication by one
+                    operands: list[Edge] = self._proxy[index].get_inputs_edges()
+                    for op_edge in operands:
+                        o = op_edge.source_node
+                        op = self.graph.nodes[o]
+                        if op.kind == NodeKinds.Constant and op.value == 1:
+                            self._remove_edge(op_edge)
+                            self._remove_node(index)                            
+                            break
+                    else:
+                        continue
+                    break

pcode_graph/torch_utils.py

@@ -0,0 +1,30 @@
+import torch
+
+class BinaryConverter:
+    """Helper to convert between scalar tensor and tensor of bits"""
+
+    def __init__(self, num_bits: int) -> None:
+        self.num_bits = num_bits
+        self._min_value = -(2 ** (num_bits - 1))
+        self._max_value = 2 ** (num_bits - 1) - 1
+        self._mask = 2 ** torch.arange(num_bits - 1, -1, -1, dtype=torch.long)
+
+    def value_to_bits(self, value: int) -> torch.Tensor:
+        """Converts the given signed value to bits"""
+
+        if value < self._min_value:
+            value = self._min_value
+        elif value > self._max_value:
+            value = self._max_value
+
+        return (
+            torch.tensor(value, dtype=torch.long)
+            .unsqueeze(-1)
+            .bitwise_and(self._mask)
+            .ne(0)
+            .float()
+        )
+
+    def bits_to_value(self, logits: torch.Tensor) -> torch.Tensor:
+        bits = (logits > 0).to(torch.long)
+        return torch.sum(self._mask * bits, -1)

pcode_graph/translator.py

@@ -0,0 +1,235 @@
+from dataclasses import dataclass
+from typing import Iterator, NamedTuple
+from loguru import logger
+import pypcode
+from pcode_graph.arch import Arch
+
+
+@dataclass(slots=True)
+class ChunkOps:
+    """Result of lifting a piece of binary."""
+
+    num_instructions: int
+    address: int
+    operations: list[pypcode.PcodeOp]
+
+
+@dataclass
+class Instruction:
+    """Result of disassembling."""
+
+    address: int
+    length: int
+    mnem: str
+    body: str
+
+    def __str__(self) -> str:
+        return self.dump(True)
+
+    def dump(self, print_address: bool):
+        s = ""
+        if print_address:
+            s += f"{self.address:05x} "
+        s += self.mnem
+        if self.body:
+            s += " " + self.body
+        return s
+
+
+@dataclass(slots=True)
+class LiftStats:
+    total_bytes: int = 0
+    lifted_bytes: int = 0
+    skipped_bytes: int = 0
+    truncated_bytes: int = 0
+    num_instructions: int = 0
+    resyncs: int = 0
+    translate_calls: int = 0
+
+    @property
+    def coverage(self) -> float:
+        return self.lifted_bytes / self.total_bytes if self.total_bytes else 1.0
+
+
+class ArchParams(NamedTuple):
+    #: Resynchronisation step, in bytes, after an undecodable instruction.
+    #: Must divide every valid instruction address, otherwise recovery lands
+    #: mid-instruction.
+    alignment: int
+    #: Longest single instruction encoding, in bytes.
+    max_instruction_size: int
+
+
+ARCH_PARAMETERS: dict[Arch, ArchParams] = {
+    Arch.x86_32: ArchParams(1, 15),
+    Arch.x86_64: ArchParams(1, 15),
+    Arch.arm_32: ArchParams(4, 4),
+    Arch.arm_64: ArchParams(4, 4),
+    Arch.mips_32: ArchParams(4, 4),
+    Arch.mips_64: ArchParams(4, 4),
+}
+
+
+def create_pcode_context(arch: Arch) -> pypcode.Context:
+    match arch:
+        case Arch.arm_32:
+            return pypcode.Context("ARM:LE:32:v8")
+        case Arch.arm_64:
+            return pypcode.Context("AARCH64:LE:64:v8A")
+            # return pypcode.Context("AARCH64:LE:64:AppleSilicon")
+        case Arch.x86_32:
+            return pypcode.Context("x86:LE:32:default")
+        case Arch.x86_64:
+            return pypcode.Context("x86:LE:64:default")
+        case Arch.mips_32:
+            return pypcode.Context("MIPS:BE:32:default")
+        case Arch.mips_64:
+            return pypcode.Context("MIPS:BE:64:default")
+        case _:
+            raise ValueError(f"Unknown architecture {arch}")
+
+
+def get_pcode_context_arch(context: pypcode.Context) -> Arch:
+    size = int(context.language.size)
+    processor = context.language.processor
+    match processor:
+        case "ARM":
+            if size == 32:
+                return Arch.arm_32
+        case "AARCH64":
+            if size == 64:
+                return Arch.arm_64
+        case "MIPS":
+            match size:
+                case 32:
+                    return Arch.mips_32
+                case 64:
+                    return Arch.mips_64
+        case "x86":
+            match size:
+                case 32:
+                    return Arch.x86_32
+                case 64:
+                    return Arch.x86_64
+
+    raise ValueError(f"Unsupported context with {processor=} and {size=}")
+
+
+class Translator:
+
+    def __init__(self, arch_context: pypcode.Context | Arch):
+        if isinstance(arch_context, Arch):
+            self.arch = arch_context
+            self.context = create_pcode_context(self.arch)
+        else:
+            self.context = arch_context
+            self.arch = get_pcode_context_arch(arch_context)
+
+        self.alignment, self.max_instr_size = ARCH_PARAMETERS[self.arch]
+
+    def translate(self, data: bytes, base_address: int) -> list[pypcode.PcodeOp]:
+        """Translates a piece of code into P-Code operations."""
+
+        ops = []
+        for chunk in self.iter_operations(data, base_address):
+            ops += chunk.operations
+        return ops
+
+    def iter_operations(
+        self,
+        data: bytes,
+        base_address: int,
+        window: int = 4096,
+        max_instructions: int = 1024,
+        stats: LiftStats | None = None,
+    ) -> Iterator[ChunkOps]:
+        """Translates a piece of code into P-Code, handling errors and avoiding
+        out-of-memory for big sections.
+
+        Each ChunkOps covers a contiguous range: undecodable bytes are skipped,
+        ``instruction_alignment`` at a time. ``window`` and ``max_instructions``
+        both bound the memory of a single call.
+        """
+
+        assert window > 1 and max_instructions > 1
+        if stats is None:
+            stats = LiftStats()
+        stats.total_bytes += len(data)
+
+        offset = 0
+        padded_data = data
+        if self.arch in {Arch.mips_32, Arch.mips_64}:
+            # Avoid delay slot exception
+            padded_data = data + b"\x00\x00\x00\x00"
+
+        while offset < len(data):
+            window_base = base_address + offset
+
+            try:
+                stats.translate_calls += 1
+                ops = self.context.translate(
+                    padded_data,
+                    window_base,
+                    max_instructions=max_instructions,
+                    offset=offset,
+                    max_bytes=window,
+                ).ops
+            except (IndexError, pypcode.BadDataError, pypcode.UnimplError, pypcode.LowlevelError):
+                ops = []
+
+            chunk: list[pypcode.PcodeOp] = []
+            num_instructions = 0
+            consumed_bytes = 0
+            truncated = False
+
+            for op in ops:
+                if op.opcode == pypcode.OpCode.IMARK:
+                    # A single IMARK may cover several instructions (grouped delay slots)
+                    # Do not assume the varnodes are sorted.
+                    end = max(i.offset + i.size for i in op.inputs) - base_address
+
+                    # pypcode silently zero-pads past the end of the buffer
+                    # and reports the instruction's full length, so the
+                    # last instruction of a section can be entirely
+                    # fabricated. Any instruction consuming a padding byte
+                    # necessarily ends past `size`, which makes this test
+                    # both sufficient and exact.
+                    if end > len(data):
+                        truncated = True
+                        break
+
+                    num_instructions += len(op.inputs)
+                    consumed_bytes = max(consumed_bytes, end - offset)
+
+                chunk.append(op)
+
+            if consumed_bytes > 0:
+                stats.lifted_bytes += consumed_bytes
+                stats.num_instructions += num_instructions
+                yield ChunkOps(num_instructions, window_base, chunk)
+                offset += consumed_bytes                
+
+            if truncated:
+                stats.truncated_bytes += len(data) - offset
+                return
+
+            if consumed_bytes == 0:                    
+                # Nothing decodable at `offset`: resynchronise.
+                stats.skipped_bytes += min(self.alignment, len(data) - offset)
+                stats.resyncs += 1
+                offset += self.alignment
+
+    def disassemble(self, code: bytes, base_address: int) -> Iterator[Instruction]:
+        data = self.context.disassemble(code, base_address=base_address)
+        for i in data.instructions:
+            yield Instruction(i.addr.offset, i.length, i.mnem, i.body)
+
+    def dump_asm(
+        self,
+        code: bytes,
+        base_address: int,
+        print_address: bool = True,
+    ) -> str:
+        return "\n".join(
+            i.dump(print_address) for i in self.disassemble(code, base_address)
+        )

pcode_graph/utils.py

@@ -0,0 +1,15 @@
+def unsigned_to_signed(value: int, size_bytes: int) -> int:
+    """Converts the given unsigned value into a signed one."""
+
+    assert size_bytes > 0
+
+    mask = 1 << (size_bytes * 8 - 1)
+    return (value ^ mask) - mask
+
+
+def signed_to_unsigned(value: int, size_bytes: int) -> int:
+    """Converts the given signed value into an unsigned one."""
+
+    return value + (value < 0) * (2 ** (size_bytes * 8))
+
+

pcode_graph/view.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+from loguru import logger
+from pyvis.network import Network
+
+from pcode_graph.graph import CDG, EdgeKinds, NodeKinds
+from pcode_graph.pcode import JUMP_OPCODES
+
+
+def draw_graph(
+    graph: CDG,
+    out_file: Path,
+    open_browser: bool = False,
+    print_ids: bool = False,
+):
+    """Creates a HTML representation of the given graph."""
+
+    parent_dir = out_file.parent
+    parent_dir.mkdir(parents=True, exist_ok=True)
+
+    network = Network(width="1200", height="1200", directed=True)
+
+    # Create nodes
+    for index, node in enumerate(graph.nodes):
+        style: dict
+
+        match node.kind:
+            case NodeKinds.InputRegister | NodeKinds.ReadMemory:
+                style = dict(color="#c0d0ff")  # Light Blue
+            case NodeKinds.OutputRegister | NodeKinds.WrittenMemory:
+                style = dict(color="#7070da")  # Blue
+            case NodeKinds.Constant:
+                style = dict(color="#ff8080")  # Red
+            case NodeKinds.Begin | NodeKinds.End | NodeKinds.External:
+                style = dict(color="#70d660", shape="box")
+            case NodeKinds.Operation | NodeKinds.Phi:
+                if node.opcode in JUMP_OPCODES:
+                    # Jumps in light Green
+                    style = dict(color="#ab2a", shape="box")
+                else:
+                    # Other operations in Grey
+                    style = dict(color="#d1c5c5", shape="box")
+            case _:
+                raise ValueError(node.kind)
+
+        label = str(node)
+        if print_ids:
+            label += " " + str(index)
+        style["label"] = label
+        network.add_node(index, **style)
+
+    # Create edges
+    for edge in graph.edges:
+        network.add_edge(
+                edge.source_node,
+                edge.destination_node,
+                style={"dashes": edge.kind == EdgeKinds.Control},
+                label=str(edge.operand_number) if edge.operand_number is not None else "",
+                color="#a1a5c5" if edge.kind == EdgeKinds.Data else "#70d660",
+            )
+
+    # Layout
+    network.force_atlas_2based(spring_length=75)
+
+    network.write_html(out_file.as_posix(), open_browser=open_browser)
+    logger.info(f"Written {out_file}")

pcode_graph/visitor.py

@@ -0,0 +1,103 @@
+from typing import Callable
+
+from pcode_graph.pcode import OpCodes
+
+
+class PcodeVisitor:
+    """A hierarchical visitor to P-Code operations."""
+
+    def __init__(self) -> None:
+        self._pcode_callbacks: dict[OpCodes | None, Callable] = {
+            opcode: getattr(self, "visit_" + opcode.name.lower()) for opcode in OpCodes
+        }
+        self._pcode_callbacks[None] = self.visit_nop
+
+    def visit(self, opcode: OpCodes | None, *args, **kwargs):
+        "Visitor main entry point."
+        return self._pcode_callbacks[opcode](*args, **kwargs)
+
+    def visit_nop(self, *args, **kwargs):
+        return self.visit_default(*args, **kwargs)
+
+    def visit_binary(self, *args, **kwargs):
+        "A binary logical or arithmetical operation."
+        return self.visit_default(*args, **kwargs)
+
+    def visit_unary(self, *args, **kwargs):
+        "An unary logical or arithmetical operation."
+        return self.visit_default(*args, **kwargs)
+
+    def visit_jump(self, *args, **kwargs):
+        "Any operation impacting the usual instruction flow."
+        return self.visit_default(*args, **kwargs)
+
+    def visit_default(self, *args, **kwargs):
+        "Catch-all callback"
+        raise NotImplementedError
+
+    # fmt: off
+    def visit_invalid(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_bool_and(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_bool_negate(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_bool_or(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_bool_xor(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_branch(self, *args, **kwargs): return self.visit_jump(*args, **kwargs)
+    def visit_branchind(self, *args, **kwargs): return self.visit_jump(*args, **kwargs)
+    def visit_call(self, *args, **kwargs): return self.visit_branch(*args, **kwargs)
+    def visit_callind(self, *args, **kwargs): return self.visit_branchind(*args, **kwargs)
+    def visit_callother(self, *args, **kwargs): return self.visit_default(*args, **kwargs)  # Note: callother is not a implemented as a branching instruction but is followed by one
+    def visit_cbranch(self, *args, **kwargs): return self.visit_jump(*args, **kwargs)
+    def visit_copy(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_abs(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_add(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_ceil(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_div(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_equal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_float2float(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_floor(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_int2float(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_less(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_lessequal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_mult(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_nan(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_neg(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_notequal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_round(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_sqrt(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_float_sub(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_float_trunc(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_int_2comp(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_int_add(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_and(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_carry(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_div(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_equal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_left(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_less(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_lessequal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_mult(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_negate(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_int_notequal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_or(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_rem(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_right(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_sborrow(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_scarry(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_sdiv(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_sext(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_int_sless(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_slessequal(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_srem(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_sright(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_sub(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)
+    def visit_int_xor(self, *args, **kwargs): return self.visit_binary(*args, **kwargs)    
+    def visit_int_zext(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_load(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_lzcount(self, *args, **kwargs): return self.visit_unary(*args, **kwargs)
+    def visit_piece(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_popcount(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_return(self, *args, **kwargs): return self.visit_branchind(*args, **kwargs)
+    def visit_segmentop(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_store(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    def visit_subpiece(self, *args, **kwargs): return self.visit_default(*args, **kwargs)
+    # fmt: on