diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/examples/custom_graph_layout.py | 89 | ||||
| -rw-r--r-- | python/flowgraph.py | 71 |
2 files changed, 152 insertions, 8 deletions
diff --git a/python/examples/custom_graph_layout.py b/python/examples/custom_graph_layout.py new file mode 100644 index 00000000..b76353de --- /dev/null +++ b/python/examples/custom_graph_layout.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# Copyright (c) 2015-2024 Vector 35 Inc +# +# 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 random + +from binaryninja import FlowGraph, FlowGraphLayout, FlowGraphNode + + +class RandomLayout(FlowGraphLayout): + def layout(self, graph: FlowGraph, nodes: list[FlowGraphNode]): + min_x = 0 + max_x = 0 + min_y = 0 + max_y = 0 + + max_extent = len(nodes) * 50 + + # Place nodes + for node in nodes: + x = random.randint(0, max_extent) + y = random.randint(0, max_extent) + + min_x = min(x, min_x) + max_x = max(x + node.width, max_x) + + min_y = min(y, min_y) + max_y = max(y + node.height, max_y) + + node.x = x + node.y = y + + # Place edges + for node in nodes: + for edge_num,edge in enumerate(node.outgoing_edges): + points = [ + (node.x + node.width/2, node.y + node.height), + (edge.target.x + edge.target.width/2, edge.target.y + edge.target.height) + ] + node.set_outgoing_edge_points(edge_num, points) + + # Calculate graph size and node visibility + for node in nodes: + min_node_x = node.x + max_node_x = node.x + node.width + + min_node_y = node.y + max_node_y = node.y + node.height + for edge in node.outgoing_edges: + for point in edge.points: + px, py = point + min_x = min(min_x, px) + min_y = min(min_y, py) + + max_x = max(max_x, px + 1) + max_y = max(max_y, py + 1) + + min_node_x = min(min_node_x, px) + max_node_x = max(max_node_x, px+1) + + min_node_y = min(min_node_y, py) + max_node_y = max(max_node_y, py+1) + + node.set_visibility_region(int(min_node_x), int(min_node_y), int(max_node_x - min_node_x), int(max_node_y - min_node_y)) + + graph.width = int(max_x - min_x) + graph.horizontal_block_margin*2 + graph.height = int(max_y - min_y) + graph.vertical_block_margin*2 + + return True + +layout = RandomLayout() +layout.register("Random Layout") diff --git a/python/flowgraph.py b/python/flowgraph.py index 431b8657..c26b20a7 100644 --- a/python/flowgraph.py +++ b/python/flowgraph.py @@ -21,7 +21,7 @@ import ctypes import threading import traceback -from typing import Optional +from typing import List, Optional, Tuple # Binary Ninja components import binaryninja @@ -61,14 +61,14 @@ class FlowGraphEdge: class EdgeStyle: - def __init__(self, style=None, width=None, theme_color=None): + def __init__(self, style: Optional[EdgePenStyle] = None, width: Optional[int] = None, theme_color: Optional[ThemeColor] = None): self.style = style if style is not None else EdgePenStyle.SolidLine self.width = width if width is not None else 0 self.color = theme_color if theme_color is not None else ThemeColor.AddressColor def _to_core_struct(self) -> core.BNEdgeStyle: result = core.BNEdgeStyle() - result.style = self.style + result.style = int(self.style) result.width = self.width result.color = self.color return result @@ -190,14 +190,22 @@ class FlowGraphNode: @property def x(self): - """Flow graph block X (read-only)""" + """Flow graph block X""" return core.BNGetFlowGraphNodeX(self.handle) + @x.setter + def x(self, x: int): + return core.BNFlowGraphNodeSetX(self.handle, x) + @property def y(self): - """Flow graph block Y (read-only)""" + """Flow graph block Y""" return core.BNGetFlowGraphNodeY(self.handle) + @y.setter + def y(self, y: int): + return core.BNFlowGraphNodeSetY(self.handle, y) + @property def width(self): """Flow graph block width (read-only)""" @@ -263,7 +271,7 @@ class FlowGraphNode: core.BNSetFlowGraphNodeLines(self.handle, line_buf, len(lines)) @property - def outgoing_edges(self): + def outgoing_edges(self) -> List[FlowGraphEdge]: """Flow graph block list of outgoing edges (read-only)""" count = ctypes.c_ulonglong() edges = core.BNGetFlowGraphNodeOutgoingEdges(self.handle, count) @@ -346,6 +354,16 @@ class FlowGraphNode: def is_valid_for_graph(self, graph): return core.BNIsNodeValidForFlowGraph(graph.handle, self.handle) + def set_visibility_region(self, x: int, y: int, w: int, h: int): + core.BNFlowGraphNodeSetVisibilityRegion(self.handle, x, y, w, h) + + def set_outgoing_edge_points(self, edge_num: int, points: List[Tuple[float, float]]): + point_buf = (core.BNPoint * len(points))() + for i in range(0, len(points)): + point_buf[i].x = points[i][0] + point_buf[i].y = points[i][1] + core.BNFlowGraphNodeSetOutgoingEdgePoints(self.handle, edge_num, point_buf, len(points)) + class FlowGraphLayoutRequest: def __init__(self, graph, callback=None): @@ -601,14 +619,22 @@ class FlowGraph: @property def width(self): - """Flow graph width (read-only)""" + """Flow graph width""" return core.BNGetFlowGraphWidth(self.handle) + @width.setter + def width(self, width: int): + return core.BNFlowGraphSetWidth(self.handle, width) + @property def height(self): - """Flow graph height (read-only)""" + """Flow graph height""" return core.BNGetFlowGraphHeight(self.handle) + @height.setter + def height(self, height: int): + return core.BNFlowGraphSetHeight(self.handle, height) + @property def horizontal_block_margin(self): return core.BNGetHorizontalFlowGraphNodeMargin(self.handle) @@ -835,3 +861,32 @@ class CoreFlowGraph(FlowGraph): if not graph: return None return CoreFlowGraph(graph) + + +class FlowGraphLayout: + def __init__(self, handle: Optional[core.BNCustomFlowGraphLayout] = None): + if handle is not None: + self.handle = core.handle_of_type(handle, core.BNFlowGraphLayout) + + def register(self, name: str): + """ + Register a custom layout with the API + """ + self._cb = core.BNCustomFlowGraphLayout() + self._cb.context = 0 + self._cb.layout = self._cb.layout.__class__(self._layout) + self.handle = core.BNRegisterFlowGraphLayout(name, self._cb) + + def _layout(self, ctxt, graph_handle: core.BNFlowGraphHandle, node_handles: 'ctypes.pointer[core.BNFlowGraphNodeHandle]', node_handle_count: int) -> bool: + try: + graph = FlowGraph(handle=graph_handle) + nodes = [] + for i in range(node_handle_count): + nodes.append(FlowGraphNode(graph=graph, handle=node_handles[i])) + return self.layout(graph, nodes) + except: + log_error(traceback.format_exc()) + return False + + def layout(self, graph: FlowGraph, nodes: List[FlowGraphNode]) -> bool: + return False |