|
| 1 | +use binaryninja::flowgraph::edge::Point; |
| 2 | +use binaryninja::flowgraph::layout::{register_flowgraph_layout, FlowGraphLayout}; |
| 3 | +use binaryninja::flowgraph::{FlowGraph, FlowGraphNode}; |
| 4 | +use binaryninja::rc::Ref; |
| 5 | +use petgraph::prelude::EdgeRef; |
| 6 | +use std::cmp::min; |
| 7 | +use std::collections::{HashMap, HashSet}; |
| 8 | + |
| 9 | +pub struct StableGraphBuilder; |
| 10 | + |
| 11 | +impl StableGraphBuilder { |
| 12 | + pub fn new() -> Self { |
| 13 | + Self {} |
| 14 | + } |
| 15 | + |
| 16 | + pub fn build( |
| 17 | + self, |
| 18 | + nodes: &[FlowGraphNode], |
| 19 | + ) -> petgraph::stable_graph::StableDiGraph<Ref<FlowGraphNode>, ()> { |
| 20 | + let mut graph = petgraph::stable_graph::StableDiGraph::<Ref<FlowGraphNode>, ()>::new(); |
| 21 | + let mut node_idx_map = HashMap::<Ref<FlowGraphNode>, petgraph::graph::NodeIndex>::new(); |
| 22 | + for node in nodes { |
| 23 | + let owned_node = node.to_owned(); |
| 24 | + node_idx_map.insert(owned_node.clone(), graph.add_node(owned_node)); |
| 25 | + } |
| 26 | + for node in nodes { |
| 27 | + let node_idx = node_idx_map.get(node).unwrap(); |
| 28 | + for edge in &node.outgoing_edges() { |
| 29 | + let target_node_idx = node_idx_map.get(&edge.target).unwrap(); |
| 30 | + graph.add_edge(*node_idx, *target_node_idx, ()); |
| 31 | + } |
| 32 | + } |
| 33 | + graph |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +struct SugiyamaLayout; |
| 38 | + |
| 39 | +impl FlowGraphLayout for SugiyamaLayout { |
| 40 | + fn layout(&self, graph: &FlowGraph, nodes: &[FlowGraphNode]) -> bool { |
| 41 | + let mut config = rust_sugiyama::configure::Config::default(); |
| 42 | + config.vertex_spacing = 5.0; |
| 43 | + |
| 44 | + let vertex_size = |_, node: &Ref<FlowGraphNode>| { |
| 45 | + let (width, height) = node.size(); |
| 46 | + (width as f64 * 1.2, height as f64) |
| 47 | + }; |
| 48 | + let pet_graph = StableGraphBuilder::new().build(nodes); |
| 49 | + let layouts = rust_sugiyama::from_graph(&pet_graph, &vertex_size, &config); |
| 50 | + |
| 51 | + // Position graph nodes |
| 52 | + for (nodes, _, _) in &layouts { |
| 53 | + for (node_idx, (x, y)) in nodes { |
| 54 | + let node = pet_graph.node_weight(*node_idx).unwrap(); |
| 55 | + node.set_position(*x as i32, *y as i32); |
| 56 | + } |
| 57 | + } |
| 58 | + |
| 59 | + // Add edges to graph nodes |
| 60 | + for (nodes, _, _) in &layouts { |
| 61 | + for (node_idx, (x, y)) in nodes { |
| 62 | + let node = pet_graph.node_weight(*node_idx).unwrap(); |
| 63 | + let (width, height) = node.size(); |
| 64 | + for (edge_idx, edge) in node.outgoing_edges().iter().enumerate() { |
| 65 | + let from_point_x = x + (width as f64 / 2.0); |
| 66 | + let from_point_y = y + height as f64; |
| 67 | + let from_point = Point { |
| 68 | + x: from_point_x as f32, |
| 69 | + y: from_point_y as f32, |
| 70 | + }; |
| 71 | + let (target_node_x, target_node_y) = edge.target.position(); |
| 72 | + let (target_node_width, _) = edge.target.size(); |
| 73 | + let to_point_x = target_node_x as f64 + (target_node_width as f64 / 2.0); |
| 74 | + let to_point_y = target_node_y; |
| 75 | + let to_point = Point { |
| 76 | + x: to_point_x as f32, |
| 77 | + y: to_point_y as f32, |
| 78 | + }; |
| 79 | + node.set_outgoing_edge_points(edge_idx, &[from_point, to_point]); |
| 80 | + } |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + // Calculate graph size and node visibility |
| 85 | + let mut min_x = f32::MAX; |
| 86 | + let mut min_y = f32::MAX; |
| 87 | + let mut max_x = f32::MIN; |
| 88 | + let mut max_y = f32::MIN; |
| 89 | + |
| 90 | + for node in nodes { |
| 91 | + let (node_x, node_y) = node.position(); |
| 92 | + let (node_width, node_height) = node.size(); |
| 93 | + |
| 94 | + // Initialize per-node bounds based on the node's current box |
| 95 | + let mut min_node_x = node_x; |
| 96 | + let mut max_node_x = node_x + node_width; |
| 97 | + let mut min_node_y = node_y; |
| 98 | + let mut max_node_y = node_y + node_height; |
| 99 | + |
| 100 | + for edge in &node.outgoing_edges() { |
| 101 | + for point in &edge.points { |
| 102 | + let px = point.x; |
| 103 | + let py = point.y; |
| 104 | + |
| 105 | + // Update Global Graph Bounds |
| 106 | + min_x = min_x.min(px); |
| 107 | + min_y = min_y.min(py); |
| 108 | + max_x = max_x.max(px + 1.0); |
| 109 | + max_y = max_y.max(py + 1.0); |
| 110 | + |
| 111 | + // Update Node Visibility Bounds |
| 112 | + min_node_x = min_node_x.min(px as i32); |
| 113 | + max_node_x = max_node_x.max(px as i32 + 1); |
| 114 | + min_node_y = min_node_y.min(py as i32); |
| 115 | + max_node_y = max_node_y.max(py as i32 + 1); |
| 116 | + } |
| 117 | + } |
| 118 | + |
| 119 | + node.set_visibility_region( |
| 120 | + min_node_x, |
| 121 | + min_node_y, |
| 122 | + (max_node_x - min_node_x), |
| 123 | + (max_node_y - min_node_y), |
| 124 | + ); |
| 125 | + } |
| 126 | + |
| 127 | + // Set final graph dimensions |
| 128 | + if min_x != f32::MAX { |
| 129 | + let (horiz_node_margin, vert_node_margin) = graph.node_margins(); |
| 130 | + let final_graph_width = (max_x - min_x) as i32 + horiz_node_margin * 2; |
| 131 | + let final_graph_height = (max_y - min_y) as i32 + vert_node_margin * 2; |
| 132 | + graph.set_size(final_graph_width, final_graph_height); |
| 133 | + } |
| 134 | + |
| 135 | + true |
| 136 | + } |
| 137 | +} |
| 138 | + |
| 139 | +/// # Safety |
| 140 | +/// This function is called from Binary Ninja once to initialize the plugin. |
| 141 | +#[allow(non_snake_case)] |
| 142 | +#[unsafe(no_mangle)] |
| 143 | +pub unsafe extern "C" fn CorePluginInit() -> bool { |
| 144 | + // Register flow graph layout |
| 145 | + register_flowgraph_layout("Sugiyama", SugiyamaLayout); |
| 146 | + true |
| 147 | +} |
0 commit comments