pcode-graph 1.0.0__py3-none-any.whl

pcode_graph/cli.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""
+Dumps Control and Data Graph (CDG) and various data from a binary or assembly file.
+
+Usage:
+    cdg ( asm | functions | html | md | pcode | table ) <input_path> [options]
+
+Dump commands:
+    asm        Assembly dump
+    functions  Functions from symbol table (using LIEF, binary must not be stripped)
+    html       CDG in HTML format
+    md         CDG in markdown/mermaid format
+    pcode      P-Code operations
+    table      Markdown table of the P-Code enriched by static analysis results
+
+Options:
+    -f, --function <name>       Take only the function with given symbol name.
+    -c, --chunk <begin:end>     Pass "0xbeaf:0xdead" to lift only the chunk started at 0xbeaf and ending before 0xdead.
+    -a, --arch <arch>           Use the given architecture (x86|arm|mips)_(64|32) or auto [Default: auto].
+    -b, --base-address <int>    Consider the given assembly starts at given address [Default: 0].
+    -D, --dataflow-only         Only extract dataflow graph (and no control flow).
+    -o, --output <path>         Where to write the output.
+    -h, --help                  This help.
+    -d, --debug                 Debug mode: do no catch exceptions.    
+"""
+
+from pathlib import Path
+import sys
+from docopt import docopt
+from lief import Binary
+import lief
+from loguru import logger
+import pypcode
+from pcode_graph.arch import Arch
+from pcode_graph.asm import Assembler
+from pcode_graph.lief_importer import (
+    get_function,
+    iter_code_sections,
+    iter_functions,
+    get_architecture,
+    lookup_chunk,
+    parse_binary,
+)
+from pcode_graph.log import setup_logger
+from pcode_graph.translator import Translator
+from pcode_graph.pcode import dump_operation
+from pcode_graph.analysis import RichPcodeList
+from pcode_graph.maker import MakerFlags, make_graph
+from pcode_graph.view import draw_graph
+from magika import Magika
+
+
+def main():
+    args = docopt(__doc__)
+
+    try:
+        if args["--debug"]:
+            setup_logger(stderr_verbosity="DEBUG")
+        else:
+            lief.logging.disable()
+            setup_logger()
+
+        input_path = Path(args["<input_path>"])
+        assert input_path.exists()
+
+        binary: Binary | None = None
+        assert input_path.stat().st_size > 0, "Input file is empty"
+        if not Magika().identify_path(input_path).output.is_text:
+            binary = parse_binary(input_path)
+
+        if args["--arch"] != "auto":
+            arch: Arch = Arch(args["--arch"])
+        else:
+            assert binary, "Assembly file requires --arch"
+            arch = get_architecture(binary)
+
+        output_path = args["--output"]
+        if output_path:
+            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+            output = open(output_path, "w")
+        else:
+            output = sys.stdout
+
+        if args["functions"]:
+            assert binary, "Command 'functions' requires a binary file"
+            for f in iter_functions(binary):
+                print(
+                    f"{hex(f.address)}: {f.name}: {len(f.content)} bytes", file=output
+                )
+            sys.exit(0)
+
+        translator = Translator(arch)
+
+        function_name = args["--function"]
+        chunk_begin_end = args["--chunk"]
+        operations: list[pypcode.PcodeOp] = []
+
+        if binary:
+            if function_name or chunk_begin_end:
+                if function_name:
+                    # Lift function
+                    function = get_function(binary, function_name)
+                    start_address = function.address
+                    code = function.content
+                else:
+                    # Lift piece of code
+                    start_address, end_address = [
+                        int(a, 16) for a in chunk_begin_end.split(":")
+                    ]
+                    code = lookup_chunk(binary, start_address, end_address)
+                if args["asm"]:
+                    print(translator.dump_asm(code, start_address, False), file=output)
+                    sys.exit(0)
+                operations = translator.translate(code, start_address)
+            else:
+                # Lift everything
+                for section in iter_code_sections(binary):
+                    if args["asm"]:
+                        print(f"; Section at {hex(section.address)}", file=output)
+                        print(
+                            translator.dump_asm(
+                                section.content, section.address, False
+                            ),
+                            file=output,
+                        )
+                    else:
+                        operations += translator.translate(
+                            section.content, section.address
+                        )
+                if args["asm"]:
+                    sys.exit(0)
+        else:
+            assert (
+                not function_name
+            ), "Option --function not supported for assembly files."
+            assert (
+                not chunk_begin_end
+            ), "Option --chunk not supported for assembly files."
+            assembly = input_path.read_text()
+            assembler = Assembler(arch)
+            start_address = int(args["--base-address"])
+            code = assembler.assemble(assembly, start_address)
+            if args["asm"]:
+                print(translator.dump_asm(code, start_address, False), file=output)
+                sys.exit(0)
+            operations = translator.translate(code, start_address)
+
+        if args["pcode"]:
+            for op in operations:
+                print(dump_operation(op), file=output)
+            sys.exit(0)
+
+        logger.info(f"Analyze {len(operations)} P-Code operations")
+        rpl = RichPcodeList(arch, operations)
+
+        if args["table"]:
+            print(rpl.dump())
+            sys.exit(0)
+
+        flags = MakerFlags(build_cfg=not args["--dataflow-only"])
+        cdg = make_graph(rpl, flags)
+
+        if args["md"]:
+            print(str(cdg), file=output)
+            sys.exit(0)
+
+        assert output_path, "HTML output file is mandatory"
+        output.close()
+        draw_graph(cdg, Path(output_path))
+
+    except Exception as e:
+        if args["--debug"]:
+            raise
+        print(str(e), file=sys.stderr)
+        sys.exit(1)

pcode_graph/emulator.py

@@ -0,0 +1,419 @@
+from typing import override
+from pcode_graph.graph import (
+    CDG,
+    EdgeKinds,
+    Node,
+    NodeIndex,
+    NodeKinds,
+    OperandNumber,
+    RegisterId,
+    Value,
+)
+from pcode_graph.utils import signed_to_unsigned, unsigned_to_signed
+from pcode_graph.visitor import PcodeVisitor
+
+
+class Emulator(PcodeVisitor):
+    """Emulate the execution of an homogeneous graph"""
+
+    def __init__(self, graph: CDG) -> None:
+        super().__init__()
+        self.graph = graph
+
+        # Known value of every node
+        self.values: dict[NodeIndex, int] = {}
+
+        # Inputs and outputs node indexes sorted by register id
+        self.inputs: dict[RegisterId, NodeIndex] = {}
+        self.outputs: dict[RegisterId, NodeIndex] = {}
+
+        for index, node in enumerate(self.graph.nodes):
+            if node.kind == NodeKinds.InputRegister:
+                self.inputs[node.register_id] = index
+            elif node.kind == NodeKinds.OutputRegister:
+                self.outputs[node.register_id] = index
+            elif node.kind == NodeKinds.Constant:
+                self.values[index] = node.value
+
+    def clear(self):
+        self.values.clear()
+
+    def set_input(self, index: NodeIndex, value: Value):
+        self.values[index] = value
+
+    def set_input_register(self, register_id: RegisterId, value: Value):
+        index = self.inputs[register_id]
+        self.set_input(index, value)
+
+    def set_input_registers(self, values_by_reg_id: dict[RegisterId, Value]):
+        for rid, value in values_by_reg_id.items():
+            self.set_input_register(rid, value)
+
+    def emulate(self):
+
+        # TODO: support CFG and Phi nodes
+        proxy = self.graph.compute_edge_proxy()
+        propagated = True
+        while propagated:
+            propagated = False
+
+            for index, node in enumerate(self.graph.nodes):
+                if index in self.values:
+                    # Already computed
+                    continue
+                p = proxy[index]
+                match node.kind:
+
+                    case NodeKinds.Operation:
+                        operands: dict[
+                            OperandNumber | None, tuple[NodeIndex, Value]
+                        ] = {}
+
+                        for in_edge in p.in_edges:
+                            if in_edge.kind != EdgeKinds.Data:
+                                continue
+
+                            value = self.values.get(in_edge.source_node)
+                            if value is None:
+                                break
+
+                            assert in_edge.operand_number not in operands
+                            operands[in_edge.operand_number] = (
+                                in_edge.source_node,
+                                value,
+                            )
+                        else:
+                            # We can evaluate this node
+                            self.values[index] = self.visit(node.opcode, node, operands)
+                            propagated = True
+
+                    case NodeKinds.OutputRegister:
+                        (in_edge,) = p.get_inputs_edges()
+                        value = self.values.get(in_edge.source_node)
+                        if value is not None:
+                            self.values[index] = value
+                            # Do not set propagated as OutputRegister nodes are never used as inputs
+
+    def get_output(self, register_id) -> Value:
+        index = self.outputs[register_id]
+        return self.values[index]
+
+    def get_outputs(self) -> dict[RegisterId, Value]:
+        out_values = {}
+        for register_id, index in self.outputs.items():
+            out_values[register_id] = self.values[index]
+        return out_values
+
+    @override
+    def visit_bool_and(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int((i1 != 0) and (i2 != 0))
+
+    @override
+    def visit_bool_negate(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        return int(values[0] == 0)
+
+    @override
+    def visit_bool_or(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int((i1 != 0) or (i2 != 0))
+
+    @override
+    def visit_bool_xor(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int((i1 != 0) ^ (i2 != 0))
+
+    @override
+    def visit_int_2comp(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i = values[0][1]
+        mask = 2**node.size_bits - 1
+        return (mask ^ i) + 1
+
+    @override
+    def visit_int_add(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 + i2
+
+    @override
+    def visit_int_and(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 & i2
+
+    @override
+    def visit_int_carry(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int((i1 + i2) >= 2**node.size_bits)
+
+    @override
+    def visit_int_div(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        if i2 == 0:
+            return 0
+        return i1 // i2
+
+    @override
+    def visit_int_equal(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int(i1 == i2)
+
+    @override
+    def visit_int_left(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 << i2
+
+    @override
+    def visit_int_less(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int(i1 < i2)
+
+    @override
+    def visit_int_lessequal(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int(i1 <= i2)
+
+    @override
+    def visit_int_mult(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 * i2
+
+    @override
+    def visit_int_negate(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i = values[0][1]
+        mask = 2 ** (node.size * 8) - 1
+        return mask ^ i
+
+    @override
+    def visit_int_notequal(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return int(i1 != i2)
+
+    @override
+    def visit_int_or(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 | i2
+
+    @override
+    def visit_int_rem(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        if i2 == 0:
+            return 0
+        q = i1 // i2
+        return i1 - q * i2
+
+    @override
+    def visit_int_right(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 >> i2
+
+    @override
+    def visit_int_sborrow(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        input_size = self.graph.nodes[values[0][0]].size
+        max = 2 ** (input_size * 8 - 1)
+        return int(
+            not (
+                -max
+                <= (
+                    unsigned_to_signed(i1, input_size)
+                    - unsigned_to_signed(i2, input_size)
+                    < max
+                )
+            )
+        )
+
+    @override
+    def visit_int_scarry(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        input_size = self.graph.nodes[values[0][0]].size
+        max = 2 ** (input_size * 8 - 1)
+        return int(
+            not (
+                -max
+                <= (
+                    unsigned_to_signed(i1, input_size)
+                    + unsigned_to_signed(i2, input_size)
+                    < max
+                )
+            )
+        )
+
+    @override
+    def visit_int_sdiv(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        if i2 == 0:
+            return 0
+        s1, s2 = unsigned_to_signed(i1, node.size), unsigned_to_signed(i2, node.size)
+        return signed_to_unsigned(int(s1 / s2), node.size)
+
+    @override
+    def visit_int_sext(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i = values[0][1]
+        input_size = self.graph.nodes[values[0][0]].size
+        s = unsigned_to_signed(i, input_size)
+        return signed_to_unsigned(s, node.size)
+
+    @override
+    def visit_int_sless(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        input_size = self.graph.nodes[values[0][0]].size
+        return int(
+            unsigned_to_signed(i1, input_size) < unsigned_to_signed(i2, input_size)
+        )
+
+    @override
+    def visit_int_slessequal(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        input_size = self.graph.nodes[values[0][0]].size
+        return int(
+            unsigned_to_signed(i1, input_size) <= unsigned_to_signed(i2, input_size)
+        )
+
+    @override
+    def visit_int_srem(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        if i2 == 0:
+            return 0
+        s1, s2 = unsigned_to_signed(i1, node.size), unsigned_to_signed(i2, node.size)
+        q = int(s1 / s2)
+        return signed_to_unsigned(s1 - q * s2, node.size)
+
+    @override
+    def visit_int_sright(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        s1, s2 = unsigned_to_signed(i1, node.size), unsigned_to_signed(i2, node.size)
+        return signed_to_unsigned(s1 >> s2, node.size)
+
+    @override
+    def visit_int_sub(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 - i2
+
+    @override
+    def visit_int_xor(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i1, i2 = values[0][1], values[1][1]
+        return i1 ^ i2
+
+    @override
+    def visit_int_zext(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        return values[0][1]
+
+    @override
+    def visit_load(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> None:
+        raise NotImplementedError
+
+    @override
+    def visit_lzcount(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        value = values[0][1]
+        max_bits = node.size * 8
+        top_bit = 1 << (max_bits - 1)
+        count = 0
+        while not (value & top_bit):
+            count += 1
+            value <<= 1
+        return count
+
+    @override
+    def visit_piece(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        # Consider the virtual machine as big endian
+        i1, i2 = values[0][1], values[1][1]
+        return (i1 << self.graph.nodes[values[1][0]].size_bits) | i2
+
+    @override
+    def visit_popcount(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        i = values[0][1]
+        return i.bit_count()
+
+    @override
+    def visit_segmentop(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        # TODO: opcode not documented in Pcode spec
+        raise NotImplementedError
+
+    @override
+    def visit_store(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> None:
+        raise NotImplementedError
+
+    @override
+    def visit_subpiece(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        # Consider the machine as big endian (see piece)
+        i1, i2 = values[0][1], values[1][1]
+        return (i1 >> (i2 * 8)) & (2**node.size_bits - 1)
+
+    @override
+    def visit_copy(
+        self, node: Node, values: dict[OperandNumber, tuple[NodeIndex, Value]]
+    ) -> int:
+        return values[0][1]