pcode-graph 1.0.0__py3-none-any.whl

pcode_graph/lief_importer.py

@@ -0,0 +1,299 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterator
+from loguru import logger
+from pypcode import OpCode
+from pcode_graph.arch import Arch
+import lief
+from pcode_graph.translator import ChunkOps, Translator
+
+lief.disable_leak_warning()
+
+
+@dataclass
+class ChunkCode:
+    address: int
+    content: bytes
+
+
+@dataclass
+class FunctionCode:
+    name: str
+    address: int
+    content: bytes
+
+
+@dataclass
+class FunctionOps(ChunkOps):
+    function_name: str
+
+
+def parse_binary(binary_path: Path) -> lief.Binary:
+    assert binary_path.exists(), f"{binary_path} not found"
+    logger.debug(f"Parse {binary_path}")
+    b = lief.parse(binary_path.as_posix())
+    assert isinstance(b, lief.Binary)
+    return b
+
+
+def get_architecture(binary: Path | lief.Binary) -> Arch:
+    """Guesses the architecture of given binary."""
+
+    if isinstance(binary, Path):
+        binary = parse_binary(binary)
+
+    # ELF
+    if isinstance(binary, lief.ELF.Binary):
+
+        match binary.header.machine_type:
+            case lief.ELF.ARCH.ARM:
+                arch = Arch.arm_32
+            case lief.ELF.ARCH.AARCH64:
+                arch = Arch.arm_64
+            case lief.ELF.ARCH.X86_64:
+                arch = Arch.x86_64
+            case lief.ELF.ARCH.I386:
+                arch = Arch.x86_32
+            case lief.ELF.ARCH.MIPS:
+                if binary.header.identity_data != lief.ELF.Header.ELF_DATA.MSB:
+                    raise ValueError("Unsupported edianness")
+                if binary.header.identity_class == lief.ELF.Header.CLASS.ELF64:
+                    arch = Arch.mips_64
+                else:
+                    flags = set(binary.header.flags_list)
+                    _P = lief.ELF.PROCESSOR_FLAGS
+                    if _P.MIPS_ABI2 in flags or bool(
+                        flags
+                        & {
+                            _P.MIPS_ARCH_3, _P.MIPS_ARCH_4, _P.MIPS_ARCH_5, _P.MIPS_ARCH_64, _P.MIPS_ARCH_64R2, _P.MIPS_ARCH_64R6, # type: ignore  # fmt: skip
+                        }
+                    ):
+                        arch = Arch.mips_64
+                    else:
+                        arch = Arch.mips_32
+            case _ as v:
+                raise ValueError(f"Unsupported CPU type {v}")
+
+    # Mach-O
+    elif isinstance(binary, lief.MachO.FatBinary):
+        arch = None
+        for sub_binary in binary:
+            match sub_binary.header.cpu_type:
+                case lief.MachO.Header.CPU_TYPE.ARM64:
+                    arch = Arch.arm_64
+                case lief.MachO.Header.CPU_TYPE.X86_64:
+                    arch = Arch.x86_64
+                case _:
+                    continue
+            break
+        else:
+            raise ValueError(f"No sub binary found with supported architecture")
+
+    elif isinstance(binary, lief.MachO.Binary):
+        match binary.header.cpu_type:
+            case lief.MachO.Header.CPU_TYPE.ARM64:
+                arch = Arch.arm_64
+            case lief.MachO.Header.CPU_TYPE.X86_64:
+                arch = Arch.x86_64
+            case _ as v:
+                raise ValueError(f"Unsupported CPU type {v}")
+
+    # PE
+    elif isinstance(binary, lief.PE.Binary):
+        match binary.header.machine:
+            case lief.PE.Header.MACHINE_TYPES.ARM64:
+                arch = Arch.arm_64
+            case lief.PE.Header.MACHINE_TYPES.AMD64:
+                arch = Arch.x86_64
+            case _ as v:
+                raise ValueError(f"Unsupported CPU type {v}")
+
+    else:
+        raise ValueError("Unsupported binary format")
+
+    return arch
+
+
+def iter_code_sections(binary: Path | lief.Binary) -> Iterator[ChunkCode]:
+
+    if isinstance(binary, Path):
+        binary = parse_binary(binary)
+
+    if binary.format == lief.Binary.FORMATS.MACHO:
+        flags = {
+            lief.MachO.Section.FLAGS.PURE_INSTRUCTIONS,  # type: ignore
+            lief.MachO.Section.FLAGS.SOME_INSTRUCTIONS,  # type: ignore
+        }
+        for section in binary.sections:
+            if flags.intersection(section.flags_list):  # type: ignore
+                yield ChunkCode(section.virtual_address, section.content.tobytes())
+
+    elif binary.format == lief.Binary.FORMATS.PE:
+        for section in binary.sections:
+            if section.has_characteristic(  # type: ignore
+                lief.PE.Section.CHARACTERISTICS.CNT_CODE
+            ) and section.has_characteristic(  # type: ignore
+                lief.PE.Section.CHARACTERISTICS.MEM_EXECUTE
+            ):
+                yield ChunkCode(section.virtual_address, section.content.tobytes())
+
+    elif binary.format == lief.Binary.FORMATS.ELF:
+        for section in binary.sections:
+            if section.has(lief.ELF.Section.FLAGS.EXECINSTR):  # type: ignore
+                yield ChunkCode(section.virtual_address, section.content.tobytes())
+
+
+def get_function(binary: Path | lief.Binary, name: str) -> FunctionCode:
+
+    if isinstance(binary, Path):
+        binary = parse_binary(binary)
+
+    for f in iter_functions(binary):
+        if f.name == name:
+            return f
+
+    raise LookupError(f"Function {name} not found in {binary}")
+
+
+def iter_functions(binary: Path | lief.Binary) -> Iterator[FunctionCode]:
+
+    if isinstance(binary, Path):
+        binary = parse_binary(binary)
+
+    cache: dict[Any, tuple[bytes, int]] = {}
+    sym: lief.ELF.Symbol
+
+    for sym in binary.symtab_symbols:  # type: ignore
+        if sym.type != lief.ELF.Symbol.TYPE.FUNC or sym.size == 0:
+            continue
+        section: lief.ELF.Section | None = sym.section
+        if section is None:
+            continue
+
+        section_id = section.name
+        if section_id not in cache:
+            cache[section_id] = (bytes(section.content), section.virtual_address)
+
+        section_content, section_address = cache[section_id]
+
+        assert isinstance(sym.name, str)
+        yield FunctionCode(
+            sym.name,
+            section_address + sym.value,
+            section_content[sym.value : sym.value + sym.size],
+        )
+
+
+def lookup_chunk(
+    binary: Path | lief.Binary,
+    start_address: int,
+    end_address: int,
+) -> bytes:
+    """Returns the code at given address from the given binary.
+
+    Args:
+        binary (Path | lief.Binary): Binary to read.
+        start_address (int): Address of the first instruction to return.
+        end_address (int): Address of the first instruction after the returned chunk.
+
+    Returns:
+        ChunkCode: A piece of binary code.
+    """
+
+    for section in iter_code_sections(binary):
+        if (
+            section.address <= start_address
+            and section.address + len(section.content) >= end_address
+        ):
+            offset = start_address - section.address
+            size = end_address - start_address
+            return section.content[offset : offset + size]
+
+    raise LookupError("Cannot find a section containing given range")
+
+
+def collect_functions(
+    binary: Path | lief.Binary,
+    translator: Translator,
+) -> Iterator[FunctionOps]:
+
+    logger.debug(f"Extract functions")
+
+    for function in iter_functions(binary):
+
+        instructions = list(translator.disassemble(function.content, function.address))
+        operations = translator.translate(function.content, function.address)
+        yield FunctionOps(
+            len(instructions),
+            function.address,
+            operations,
+            function.name,
+        )
+
+
+def collect_basic_blocks(
+    binary: Path | lief.Binary,
+    translator: Translator,
+    min_instructions: int,
+    max_instructions: int,
+) -> Iterator[ChunkOps]:
+    """Extracts the P-Code operations of basic blocks found in a binary, using LIEF library."""
+
+    logger.debug(f"Extract basic blocks")
+
+    # Extract P-Code operations of executable sections
+    for section in iter_code_sections(binary):
+
+        # Slice into basic blocks
+        num_instructions = 0
+        operations = []
+        bb_start_offset = 0
+
+        # First disassemble to translate into P-Code one instruction at a time,
+        # in order to avoid consuming too much RAM.
+        for instr in translator.disassemble(section.content, section.address):
+
+            instr_start_offset = instr.address - section.address
+            instr_end_offset = instr_start_offset + instr.length
+            instr_code = section.content[instr_start_offset:instr_end_offset]
+
+            instr_ops = translator.translate(instr_code, instr.address)
+
+            num_instructions += 1
+
+            if num_instructions <= max_instructions:
+                operations += instr_ops
+
+            last_op = instr_ops[-1]
+            if last_op.opcode not in {
+                OpCode.BRANCH,
+                OpCode.BRANCHIND,
+                OpCode.CBRANCH,
+                OpCode.RETURN,
+            }:
+                # Not at the end of a basic block
+                continue
+
+            if min_instructions <= num_instructions <= max_instructions:
+                yield ChunkOps(
+                    num_instructions,
+                    section.address + bb_start_offset,
+                    # section.content[bb_start_offset:instr_end_offset],
+                    operations,
+                )
+
+            bb_start_offset = instr_end_offset
+            num_instructions = 0
+            operations = []
+
+
+def collect_chunks(
+    binary: Path | lief.Binary,
+    translator: Translator,    
+    
+) -> Iterator[ChunkOps]:
+    """Extracts the P-Code operations of the sections found in a binary, using LIEF library."""
+
+    for section in iter_code_sections(binary):
+        yield from translator.iter_operations(section.content, section.address)        
+        

pcode_graph/log.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+import sys
+from loguru import logger
+
+
+def setup_logger(
+    log_path: Path | None = None,
+    stderr_verbosity="INFO",
+    file_verbosity="DEBUG",
+) -> int | None:
+    """Setups the global logger with two sinks: a file and stderr.
+    Returns a handle to the file sink.
+    """
+
+    logger.remove()
+    logger.add(
+        sys.stderr,
+        level=stderr_verbosity,
+        format="<level>{message}</level>",
+        backtrace=False,
+        diagnose=False,
+    )
+
+    if log_path is not None:
+        logger.info(f"Logging in {log_path}")
+        log_path.parent.mkdir(exist_ok=True, parents=True)
+        if log_path.exists():
+            logger.warning(f"Overwriting logfile: {log_path}")
+            log_path.unlink()
+        return logger.add(
+            log_path,
+            level=file_verbosity,
+            format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
+            "<level>{level: <8}</level> | "
+            "<cyan>{name}:{line}: in {function}:</cyan> <level>{message}</level>",
+        )

pcode_graph/maker.py

@@ -0,0 +1,369 @@
+from dataclasses import dataclass, field
+
+from pcode_graph.analysis import (
+    EntryIndex,
+    OpIndex,
+    RichPcodeList,
+    UnknownCodeIndex,
+    ExitIndex,
+    Var,
+    VarKinds,
+)
+from pcode_graph.asm import Assembler
+from pcode_graph.checker import check_graph
+from pcode_graph.graph import CDG, Edge, EdgeKinds, NodeIndex, NodeKinds, Node
+from pcode_graph.pcode import from_native_opcode
+from pcode_graph.registers import SPECIAL_REGISTERS
+from pcode_graph.simplifier import GraphSimplifier
+from pcode_graph.translator import Translator
+
+from pypcode import OpCode
+from loguru import logger
+
+
+@dataclass
+class MakerFlags:
+    build_cfg: bool = True
+    """Add control flow links in addition to the data flow. Defaults to True."""
+
+    ignore_reg_outputs: set[str] | None = None
+    """Do not consider the given registers as graph outputs to simplify the graphs.
+    Arch-dependent register names.
+    If None, the result of pcode.get_special_register_names() is used.
+    """
+
+
+def make_graph(prg: RichPcodeList, flags: MakerFlags | None = None) -> CDG:
+    """Makes a graph from a list of P-Code operations.
+
+    Args:
+        prg (RichPcodeList): the input list of operations, enriched by analyses.
+        flags (MakerFlags): determines what to put in the graph.
+
+    Returns:
+        CDG: the graph built.
+    """
+
+    if flags is None:
+        flags = MakerFlags()
+    ignore_reg_outputs = flags.ignore_reg_outputs
+    if ignore_reg_outputs is None:
+        ignore_reg_outputs = SPECIAL_REGISTERS[prg.arch]
+
+    # Already created input nodes, by related input variable
+    input_node_indexes: dict[Var, NodeIndex] = {}
+
+    # Already created operation nodes, by related P-Code operation
+    operation_node_indexes: dict[OpIndex, NodeIndex] = {}
+
+    # Already created phi nodes, by (sorted) possible alternatives
+    # to avoid duplicate phi nodes
+    phi_node_indexes: dict[tuple[NodeIndex, NodeIndex], NodeIndex] = {}
+
+    # Created nodes of all kinds
+    nodes: list[Node] = []
+
+    # Created edges
+    edges: list[Edge] = []
+
+    # To build the data graph recursively
+    newly_taken_ops: list[tuple[OpIndex, NodeIndex]] = []
+
+    # External code node if created
+    external_node_index: NodeIndex | None = None
+
+    # End node if created
+    end_node_index: NodeIndex | None = None
+
+    # begin node has zero index
+    if flags.build_cfg:
+        nodes.append(Node(NodeKinds.Begin))
+        begin_index = 0
+    else:
+        begin_index = -1
+
+    def get_or_create_terminal(op_index: OpIndex) -> NodeIndex:
+        nonlocal external_node_index, end_node_index
+
+        if op_index == ExitIndex:
+            if end_node_index is None:
+                end_node_index = len(nodes)
+                nodes.append(Node(NodeKinds.End))
+            return end_node_index
+        assert op_index == UnknownCodeIndex
+        if external_node_index is None:
+            external_node_index = len(nodes)
+            nodes.append(Node(NodeKinds.External))
+        return external_node_index
+
+    def get_or_create_operation(op_index: OpIndex) -> NodeIndex:
+        node_index = operation_node_indexes.get(op_index)
+        if node_index is None:
+            op = prg.operations[op_index]
+            node_index = len(nodes)
+            newly_taken_ops.append((op_index, node_index))
+            nodes.append(
+                Node(
+                    NodeKinds.Operation,
+                    op.output.size if op.output is not None else 0,
+                    opcode=from_native_opcode(op.opcode),
+                )
+            )
+            operation_node_indexes[op_index] = node_index
+
+            # logger.debug(f"Create node {node_index}: {nodes[-1]} from operation {op_index}: {dump_operation(op)}")
+        return node_index
+
+    def get_or_create_input(var: Var) -> NodeIndex:
+        node_index = input_node_indexes.get(var)
+        if node_index is None:
+            match var.kind:
+                case VarKinds.Constant:
+                    node = Node(NodeKinds.Constant, var.size, var.value)
+                case VarKinds.Register:
+                    node = Node(
+                        NodeKinds.InputRegister,
+                        var.size,
+                        register_id=var.value,
+                        register_name=var.pretty_name,
+                    )
+                case VarKinds.Ram:
+                    node = Node(NodeKinds.ReadMemory, var.size)
+                case VarKinds.IndirectAccess:
+                    node = Node(NodeKinds.ReadMemory, var.size)
+                case _:
+                    raise ValueError(f"Unexpected kind for input {var.kind}")
+            node_index = len(nodes)
+            nodes.append(node)
+            input_node_indexes[var] = node_index
+        return node_index
+
+    def add_control_edge(
+        from_node: NodeIndex, to_node: NodeIndex, operand_number: int | None
+    ):
+        edges.append(Edge(from_node, to_node, EdgeKinds.Control, operand_number))
+
+    def add_data_edge(
+        from_node: NodeIndex, to_node: NodeIndex, operand_number: int | None
+    ):
+        edges.append(Edge(from_node, to_node, EdgeKinds.Data, operand_number))
+
+    def link_defs(
+        var: Var,
+        def_op_indexes: set[OpIndex],
+        user_node_index: NodeIndex,
+        operand_number: int | None,
+    ):
+
+        # Collect data preds to see if we need a phi node
+        # Sort defs to remain determinist
+        pred_node_indexes = []
+        for def_op_index in sorted(list(def_op_indexes)):
+            if def_op_index == EntryIndex:
+                # Input variable
+                pred_node_indexes.append(get_or_create_input(var))
+            else:
+                pred_node_indexes.append(get_or_create_operation(def_op_index))
+
+        if len(pred_node_indexes) == 1:
+            # Simple def, direct link
+            add_data_edge(pred_node_indexes[0], user_node_index, operand_number)
+        else:
+            # Create phi node
+            assert pred_node_indexes
+            pred_node_indexes.sort()
+            key = tuple(pred_node_indexes)
+            phi_node_index = phi_node_indexes.get(key)
+            if phi_node_index is None:
+                phi_node = Node(NodeKinds.Phi, var.size)
+                phi_node_index = len(nodes)
+                nodes.append(phi_node)
+                phi_node_indexes[key] = phi_node_index
+                for pred_node_index in pred_node_indexes:
+                    add_data_edge(pred_node_index, phi_node_index, None)
+            add_data_edge(phi_node_index, user_node_index, operand_number)
+
+    # Build the data flow
+
+    # Start from writes to outputs registers and memory
+    for var, def_op_indexes in prg.get_exit_defs().items():
+        assert def_op_indexes
+        match var.kind:
+            case VarKinds.Register:
+                if var.pretty_name.lower().startswith("tmp"):
+                    # Ignore writes to temporaries
+                    continue
+                if var.pretty_name in ignore_reg_outputs:
+                    # Ignore writes to special registers
+                    continue
+                output_node = Node(
+                    NodeKinds.OutputRegister,
+                    var.size,
+                    register_id=var.value,
+                    register_name=var.pretty_name,
+                )
+
+            case VarKinds.IndirectAccess | VarKinds.Ram:
+                output_node = Node(NodeKinds.WrittenMemory, var.size)
+
+            case VarKinds.Unique:
+                # Ignore writes to temporaries
+                continue
+
+            case VarKinds.Constant:
+                raise ValueError(f"Unexpected output kind {var.kind}")
+
+        output_node_index = len(nodes)
+        nodes.append(output_node)
+        # logger.debug(f"Start from output {output_node_index} {output_node}")
+        link_defs(var, def_op_indexes, output_node_index, None)
+
+    if flags.build_cfg:
+        # Also add conditional branching or branching to unknown address
+        for op_index, op in enumerate(prg.operations):
+            if not prg.reachable[op_index]:
+                continue
+            if op.opcode == OpCode.CBRANCH or (
+                UnknownCodeIndex in prg.successors.get(op_index, set())
+                and op.opcode != OpCode.IMARK
+            ):
+                get_or_create_operation(op_index)
+
+    # Follow predecessors recursively until the whole dataflow graph is built
+    while newly_taken_ops:
+        op_index, node_index = newly_taken_ops.pop()
+
+        for number, (var, defs) in enumerate(prg.iter_input_defs(op_index)):
+
+            if not defs:
+                # No definition means this is a constant
+                pred_node_index = get_or_create_input(var)
+                add_data_edge(pred_node_index, node_index, number)
+                continue
+
+            link_defs(var, defs, node_index, number)
+
+    if flags.build_cfg:
+        # Build the control flow graph, jumping over nodes which
+        # do not already belong to the data flow graph.
+
+        for op_index, node_index in operation_node_indexes.items():
+
+            if not prg.reachable[op_index]:
+                continue
+
+            successors = prg.successors.get(op_index, [])
+
+            for number, successor_op_index in enumerate(successors):
+                if successor_op_index < 0:  # ie. in {UnknownCodeIndex, ExitIndex}
+                    add_control_edge(
+                        node_index, get_or_create_terminal(successor_op_index), number
+                    )
+                    continue
+
+                successor_node_index = operation_node_indexes.get(successor_op_index)
+                seen_ops: set[OpIndex] = set()
+
+                while successor_node_index is None:
+
+                    # This successor is not part of the dataflow graph, skip it.
+
+                    # Having more than one successor would imply being a CBRANCH,
+                    # but CBRANCH are always put in the dataflow, see above.
+                    (successor_op_index,) = prg.successors[successor_op_index]
+
+                    if successor_op_index in seen_ops:
+                        logger.warning("Looping CFG without side effects")
+                        add_control_edge(node_index, node_index, number)
+                        break
+
+                    seen_ops.add(successor_op_index)
+
+                    if successor_op_index < 0:  # ie. in {UnknownCodeIndex, ExitIndex}
+                        successor_node_index = get_or_create_terminal(
+                            successor_op_index
+                        )
+                        continue
+
+                    successor_node_index = operation_node_indexes.get(
+                        successor_op_index
+                    )
+
+                else:
+                    assert successor_node_index >= 0
+                    add_control_edge(node_index, successor_node_index, number)
+
+        if len(nodes) == 1:
+            # Empty graph
+            add_control_edge(begin_index, get_or_create_terminal(ExitIndex), None)
+        else:
+            # Link the entry node to the first taken operation
+            # TODO: specify entry operation
+            entry_op_index = 0
+            while entry_op_index not in operation_node_indexes:
+                successor_op_indexes = tuple(prg.successors.get(entry_op_index, ()))
+                assert len(successor_op_indexes) == 1
+                # Advance to next operation
+                entry_op_index = successor_op_indexes[0]
+            add_control_edge(begin_index, operation_node_indexes[entry_op_index], None)
+
+    graph = CDG(nodes, edges)
+
+    # Uncomment to locate a graph structure bug
+    # Comment to speed up
+    # check_graph(graph)
+
+    GraphSimplifier(graph).process()
+    check_graph(graph)
+
+    return graph
+
+
+def make_graph_from_binary(
+    translator: Translator,
+    code: bytes,
+    base_address: int,
+    flags: MakerFlags | None = None,
+) -> CDG:
+    """Makes a CDG graph from a piece of binary code.
+
+    Args:
+        translator (Translator): The P-Code translator context.
+        code (bytes): The binary code to translate.
+        base_address (int): The address of the first instruction of the given code.
+        flags (MakerFlags): determines what to put in the graph.
+
+    Returns:
+        CDG: the graph built.
+    """
+
+    operations = translator.translate(code, base_address)
+    prg = RichPcodeList(translator.arch, operations=operations)
+    return make_graph(prg, flags)
+
+
+def make_graph_from_asm(
+    assembler: Assembler,
+    translator: Translator,
+    asm: str,
+    base_address: int = 0,
+    flags: MakerFlags | None = None,
+) -> CDG:
+    """Makes a CDG graph from a piece of binary code. Mainly for testing purpose.
+
+    Args:
+        translator (Translator): The P-Code translator context.
+        asm (str): The assembly code to translate.
+        base_address (int): The address of the first instruction of the given code.
+        flags (MakerFlags): determines what to put in the graph.
+
+    Returns:
+        CDG: the graph built.
+    """
+
+    return make_graph_from_binary(
+        code=assembler.assemble(asm, base_address),
+        base_address=base_address,
+        translator=translator,
+        flags=flags,
+    )