diff options
Diffstat (limited to 'python/flowgraph.py')
| -rw-r--r-- | python/flowgraph.py | 465 |
1 files changed, 465 insertions, 0 deletions
diff --git a/python/flowgraph.py b/python/flowgraph.py new file mode 100644 index 00000000..1568bc77 --- /dev/null +++ b/python/flowgraph.py @@ -0,0 +1,465 @@ +# Copyright (c) 2018 Vector 35 LLC +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to +# deal in the Software without restriction, including without limitation the +# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or +# sell copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS +# IN THE SOFTWARE. + +import ctypes +import threading +import traceback + +# Binary Ninja components +import _binaryninjacore as core +from enums import (BranchType, InstructionTextTokenType, HighlightColorStyle, HighlightStandardColor) +import function +import binaryview +import lowlevelil +import mediumlevelil +import basicblock +import architecture +import log +import interaction +import highlight + + +class FlowGraphEdge(object): + def __init__(self, branch_type, source, target, points, back_edge): + self.type = BranchType(branch_type) + self.source = source + self.target = target + self.points = points + self.back_edge = back_edge + + def __repr__(self): + return "<%s: %s>" % (self.type.name, repr(self.target)) + + +class FlowGraphNode(object): + def __init__(self, graph, handle = None): + if handle is None: + handle = core.BNCreateFlowGraphNode(graph.handle) + self.handle = handle + self.graph = graph + + def __del__(self): + core.BNFreeFlowGraphNode(self.handle) + + def __eq__(self, value): + if not isinstance(value, FlowGraphNode): + return False + return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents) + + def __ne__(self, value): + if not isinstance(value, FlowGraphNode): + return True + return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents) + + @property + def basic_block(self): + """Basic block associated with this part of the flow graph (read-only)""" + block = core.BNGetFlowGraphBasicBlock(self.handle) + if not block: + return None + func_handle = core.BNGetBasicBlockFunction(block) + if not func_handle: + core.BNFreeBasicBlock(block) + return None + + view = binaryview.BinaryView(handle = core.BNGetFunctionData(func_handle)) + func = function.Function(view, func_handle) + + if core.BNIsLowLevelILBasicBlock(block): + block = lowlevelil.LowLevelILBasicBlock(view, block, + lowlevelil.LowLevelILFunction(func.arch, core.BNGetBasicBlockLowLevelILFunction(block), func)) + elif core.BNIsMediumLevelILBasicBlock(block): + block = mediumlevelil.MediumLevelILBasicBlock(view, block, + mediumlevelil.MediumLevelILFunction(func.arch, core.BNGetBasicBlockMediumLevelILFunction(block), func)) + else: + block = basicblock.BasicBlock(view, block) + return block + + @property + def x(self): + """Flow graph block X (read-only)""" + return core.BNGetFlowGraphNodeX(self.handle) + + @property + def y(self): + """Flow graph block Y (read-only)""" + return core.BNGetFlowGraphNodeY(self.handle) + + @property + def width(self): + """Flow graph block width (read-only)""" + return core.BNGetFlowGraphNodeWidth(self.handle) + + @property + def height(self): + """Flow graph block height (read-only)""" + return core.BNGetFlowGraphNodeHeight(self.handle) + + @property + def lines(self): + """Flow graph block list of lines""" + count = ctypes.c_ulonglong() + lines = core.BNGetFlowGraphNodeLines(self.handle, count) + block = self.basic_block + result = [] + for i in xrange(0, count.value): + addr = lines[i].addr + if (lines[i].instrIndex != 0xffffffffffffffff) and (block is not None) and hasattr(block, 'il_function'): + il_instr = block.il_function[lines[i].instrIndex] + else: + il_instr = None + color = highlight.HighlightColor._from_core_struct(lines[i].highlight) + tokens = [] + for j in xrange(0, lines[i].count): + token_type = InstructionTextTokenType(lines[i].tokens[j].type) + text = lines[i].tokens[j].text + value = lines[i].tokens[j].value + size = lines[i].tokens[j].size + operand = lines[i].tokens[j].operand + context = lines[i].tokens[j].context + confidence = lines[i].tokens[j].confidence + address = lines[i].tokens[j].address + tokens.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address, confidence)) + result.append(function.DisassemblyTextLine(tokens, addr, il_instr, color)) + core.BNFreeDisassemblyTextLines(lines, count.value) + return result + + @lines.setter + def lines(self, lines): + if isinstance(lines, str): + lines = lines.split('\n') + line_buf = (core.BNDisassemblyTextLine * len(lines))() + for i in xrange(0, len(lines)): + line = lines[i] + if isinstance(line, str): + line = function.DisassemblyTextLine([function.InstructionTextToken(InstructionTextTokenType.TextToken, line)]) + if not isinstance(line, function.DisassemblyTextLine): + line = function.DisassemblyTextLine(line) + if line.address is None: + if len(line.tokens) > 0: + line_buf[i].addr = line.tokens[0].address + else: + line_buf[i].addr = 0 + else: + line_buf[i].addr = line.address + if line.il_instruction is not None: + line_buf[i].instrIndex = line.il_instruction.instr_index + else: + line_buf[i].instrIndex = 0xffffffffffffffff + color = line.highlight + if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor): + raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor") + if isinstance(color, HighlightStandardColor): + color = highlight.HighlightColor(color) + line_buf[i].highlight = color._get_core_struct() + line_buf[i].count = len(line.tokens) + line_buf[i].tokens = (core.BNInstructionTextToken * len(line.tokens))() + for j in xrange(0, len(line.tokens)): + line_buf[i].tokens[j].type = line.tokens[j].type + line_buf[i].tokens[j].text = line.tokens[j].text + line_buf[i].tokens[j].value = line.tokens[j].value + line_buf[i].tokens[j].size = line.tokens[j].size + line_buf[i].tokens[j].operand = line.tokens[j].operand + line_buf[i].tokens[j].context = line.tokens[j].context + line_buf[i].tokens[j].confidence = line.tokens[j].confidence + line_buf[i].tokens[j].address = line.tokens[j].address + core.BNSetFlowGraphNodeLines(self.handle, line_buf, len(lines)) + + @property + def outgoing_edges(self): + """Flow graph block list of outgoing edges (read-only)""" + count = ctypes.c_ulonglong() + edges = core.BNGetFlowGraphNodeOutgoingEdges(self.handle, count) + result = [] + for i in xrange(0, count.value): + branch_type = BranchType(edges[i].type) + target = edges[i].target + if target: + target = FlowGraphNode(self.graph, core.BNNewFlowGraphNodeReference(target)) + points = [] + for j in xrange(0, edges[i].pointCount): + points.append((edges[i].points[j].x, edges[i].points[j].y)) + result.append(FlowGraphEdge(branch_type, self, target, points, edges[i].backEdge)) + core.BNFreeFlowGraphNodeOutgoingEdgeList(edges, count.value) + return result + + @property + def highlight(self): + """Gets or sets the highlight color for the node + + :Example: + >>> g = FlowGraph() + >>> node = FlowGraphNode(g) + >>> node.highlight = HighlightStandardColor.BlueHighlightColor + >>> node.highlight + <color: blue> + """ + return highlight.HighlightColor._from_core_struct(core.BNGetFlowGraphNodeHighlight(self.handle)) + + @highlight.setter + def highlight(self, color): + if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor): + raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor") + if isinstance(color, HighlightStandardColor): + color = highlight.HighlightColor(color) + core.BNSetFlowGraphNodeHighlight(self.handle, color._get_core_struct()) + + def __repr__(self): + block = self.basic_block + if block: + arch = block.arch + if arch: + return "<graph node: %s@%#x-%#x>" % (arch.name, block.start, block.end) + else: + return "<graph node: %#x-%#x>" % (block.start, block.end) + return "<graph node>" + + def __iter__(self): + count = ctypes.c_ulonglong() + lines = core.BNGetFlowGraphNodeLines(self.handle, count) + block = self.basic_block + try: + for i in xrange(0, count.value): + addr = lines[i].addr + if (lines[i].instrIndex != 0xffffffffffffffff) and (block is not None) and hasattr(block, 'il_function'): + il_instr = block.il_function[lines[i].instrIndex] + else: + il_instr = None + tokens = [] + for j in xrange(0, lines[i].count): + token_type = InstructionTextTokenType(lines[i].tokens[j].type) + text = lines[i].tokens[j].text + value = lines[i].tokens[j].value + size = lines[i].tokens[j].size + operand = lines[i].tokens[j].operand + context = lines[i].tokens[j].context + confidence = lines[i].tokens[j].confidence + address = lines[i].tokens[j].address + tokens.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address, confidence)) + yield function.DisassemblyTextLine(tokens, addr, il_instr) + finally: + core.BNFreeDisassemblyTextLines(lines, count.value) + + def add_outgoing_edge(self, edge_type, target): + core.BNAddFlowGraphNodeOutgoingEdge(self.handle, edge_type, target.handle) + + +class FlowGraph(object): + def __init__(self, handle = None): + if handle is None: + handle = core.BNCreateFlowGraph() + self.handle = handle + self._on_complete = None + self._cb = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self._complete) + + def __del__(self): + self.abort() + core.BNFreeFlowGraph(self.handle) + + def __eq__(self, value): + if not isinstance(value, FlowGraph): + return False + return ctypes.addressof(self.handle.contents) == ctypes.addressof(value.handle.contents) + + def __ne__(self, value): + if not isinstance(value, FlowGraph): + return True + return ctypes.addressof(self.handle.contents) != ctypes.addressof(value.handle.contents) + + @property + def function(self): + """Function for a flow graph""" + func = core.BNGetFunctionForFlowGraph(self.handle) + if func is None: + return None + return function.Function(binaryview.BinaryView(handle = core.BNGetFunctionData(func)), func) + + @function.setter + def function(self, func): + if func is not None: + func = func.handle + core.BNSetFunctionForFlowGraph(self.handle, func) + + @property + def complete(self): + """Whether flow graph layout is complete (read-only)""" + return core.BNIsFlowGraphLayoutComplete(self.handle) + + @property + def nodes(self): + """List of nodes in graph (read-only)""" + count = ctypes.c_ulonglong() + blocks = core.BNGetFlowGraphNodes(self.handle, count) + result = [] + for i in xrange(0, count.value): + result.append(FlowGraphNode(self, core.BNNewFlowGraphNodeReference(blocks[i]))) + core.BNFreeFlowGraphNodeList(blocks, count.value) + return result + + @property + def has_nodes(self): + """Whether the flow graph has at least one node (read-only)""" + return core.BNFlowGraphHasNodes(self.handle) + + @property + def width(self): + """Flow graph width (read-only)""" + return core.BNGetFlowGraphWidth(self.handle) + + @property + def height(self): + """Flow graph height (read-only)""" + return core.BNGetFlowGraphHeight(self.handle) + + @property + def horizontal_block_margin(self): + return core.BNGetHorizontalFlowGraphBlockMargin(self.handle) + + @horizontal_block_margin.setter + def horizontal_block_margin(self, value): + core.BNSetFlowGraphBlockMargins(self.handle, value, self.vertical_block_margin) + + @property + def vertical_block_margin(self): + return core.BNGetVerticalFlowGraphBlockMargin(self.handle) + + @vertical_block_margin.setter + def vertical_block_margin(self, value): + core.BNSetFlowGraphBlockMargins(self.handle, self.horizontal_block_margin, value) + + @property + def is_il(self): + return core.BNIsILFlowGraph(self.handle) + + @property + def is_low_level_il(self): + return core.BNIsLowLevelILFlowGraph(self.handle) + + @property + def is_medium_level_il(self): + return core.BNIsMediumLevelILFlowGraph(self.handle) + + @property + def il_function(self): + if self.is_low_level_il: + il_func = core.BNGetFlowGraphLowLevelILFunction(self.handle) + if not il_func: + return None + function = self.function + if function is None: + return None + return lowlevelil.LowLevelILFunction(function.arch, il_func, function) + if self.is_medium_level_il: + il_func = core.BNGetFlowGraphMediumLevelILFunction(self.handle) + if not il_func: + return None + function = self.function + if function is None: + return None + return mediumlevelil.MediumLevelILFunction(function.arch, il_func, function) + return None + + @il_function.setter + def il_function(self, func): + if isinstance(func, lowlevelil.LowLevelILFunction): + core.BNSetFlowGraphLowLevelILFunction(self.handle, func.handle) + core.BNSetFlowGraphMediumLevelILFunction(self.handle, None) + elif isinstance(func, mediumlevelil.MediumLevelILFunction): + core.BNSetFlowGraphLowLevelILFunction(self.handle, None) + core.BNSetFlowGraphMediumLevelILFunction(self.handle, func.handle) + elif func is None: + core.BNSetFlowGraphLowLevelILFunction(self.handle, None) + core.BNSetFlowGraphMediumLevelILFunction(self.handle, None) + else: + raise TypeError("expected IL function for setting il_function property") + + def __setattr__(self, name, value): + try: + object.__setattr__(self, name, value) + except AttributeError: + raise AttributeError("attribute '%s' is read only" % name) + + def __repr__(self): + function = self.function + if function is None: + return "<flow graph>" + return "<graph of %s>" % repr(function) + + def __iter__(self): + count = ctypes.c_ulonglong() + nodes = core.BNGetFlowGraphNodes(self.handle, count) + try: + for i in xrange(0, count.value): + yield FlowGraphNode(self, core.BNNewFlowGraphNodeReference(nodes[i])) + finally: + core.BNFreeFlowGraphNodeList(nodes, count.value) + + def _complete(self, ctxt): + try: + if self._on_complete is not None: + self._on_complete() + except: + log.log_error(traceback.format_exc()) + + def layout(self): + core.BNStartFlowGraphLayout(self.handle) + + def _wait_complete(self): + self._wait_cond.acquire() + self._wait_cond.notify() + self._wait_cond.release() + + def layout_and_wait(self): + self._wait_cond = threading.Condition() + self.on_complete(self._wait_complete) + self.layout() + + self._wait_cond.acquire() + while not self.complete: + self._wait_cond.wait() + self._wait_cond.release() + + def on_complete(self, callback): + self._on_complete = callback + core.BNSetFlowGraphCompleteCallback(self.handle, None, self._cb) + + def abort(self): + core.BNAbortFlowGraph(self.handle) + + def get_nodes_in_region(self, left, top, right, bottom): + count = ctypes.c_ulonglong() + nodes = core.BNGetFlowGraphNodesInRegion(self.handle, left, top, right, bottom, count) + result = [] + for i in xrange(0, count.value): + result.append(FlowGraphNode(self, core.BNNewFlowGraphNodeReference(nodes[i]))) + core.BNFreeFlowGraphNodeList(nodes, count.value) + return result + + def append(self, node): + return core.BNAddFlowGraphNode(self.handle, node.handle) + + def __getitem__(self, i): + node = core.BNGetFlowGraphNode(self.handle, i) + if node is None: + return None + return FlowGraphNode(self, node) + + def show(self, title): + interaction.show_graph_report(title, self) |