1414
1515//! Interfaces for creating and displaying pretty CFGs in Binary Ninja.
1616
17- use binaryninjacore_sys:: * ;
18-
1917use crate :: disassembly:: DisassemblyTextLine ;
18+ use binaryninjacore_sys:: * ;
19+ use std:: slice;
2020
2121use crate :: rc:: * ;
2222
23+ use crate :: basic_block:: { BasicBlock , BlockContext } ;
24+ use crate :: function:: HighlightColor ;
25+ use crate :: render_layer:: CoreRenderLayer ;
2326use std:: marker:: PhantomData ;
2427
2528pub type BranchType = BNBranchType ;
@@ -37,21 +40,89 @@ impl FlowGraph {
3740 Self { handle : raw }
3841 }
3942
43+ pub ( crate ) unsafe fn ref_from_raw ( raw : * mut BNFlowGraph ) -> Ref < Self > {
44+ Ref :: new ( Self { handle : raw } )
45+ }
46+
4047 pub fn new ( ) -> Ref < Self > {
41- unsafe { Ref :: new ( FlowGraph :: from_raw ( BNCreateFlowGraph ( ) ) ) }
48+ unsafe { FlowGraph :: ref_from_raw ( BNCreateFlowGraph ( ) ) }
49+ }
50+
51+ pub fn nodes < ' a > ( & self ) -> Vec < Ref < FlowGraphNode < ' a > > > {
52+ let mut count: usize = 0 ;
53+ let nodes_ptr = unsafe { BNGetFlowGraphNodes ( self . handle , & mut count as * mut usize ) } ;
54+
55+ let nodes = unsafe { slice:: from_raw_parts_mut ( nodes_ptr, count) } ;
56+
57+ let mut result = vec ! [ ] ;
58+ result. reserve ( count) ;
59+
60+ for i in 0 ..count {
61+ result. push ( unsafe { RefCountable :: inc_ref ( & FlowGraphNode :: from_raw ( nodes[ i] ) ) } ) ;
62+ }
63+
64+ unsafe { BNFreeFlowGraphNodeList ( nodes_ptr, count) } ;
65+
66+ result
67+ }
68+
69+ pub fn get_node < ' a > ( & self , i : usize ) -> Option < Ref < FlowGraphNode < ' a > > > {
70+ let node_ptr = unsafe { BNGetFlowGraphNode ( self . handle , i) } ;
71+ if node_ptr. is_null ( ) {
72+ None
73+ } else {
74+ Some ( unsafe { Ref :: new ( FlowGraphNode :: from_raw ( node_ptr) ) } )
75+ }
76+ }
77+
78+ pub fn get_node_count ( & self ) -> usize {
79+ unsafe { BNGetFlowGraphNodeCount ( self . handle ) }
80+ }
81+
82+ pub fn has_nodes ( & self ) -> bool {
83+ unsafe { BNFlowGraphHasNodes ( self . handle ) }
4284 }
4385
4486 pub fn append ( & self , node : & FlowGraphNode ) -> usize {
4587 unsafe { BNAddFlowGraphNode ( self . handle , node. handle ) }
4688 }
4789
90+ pub fn replace ( & self , index : usize , node : & FlowGraphNode ) {
91+ unsafe { BNReplaceFlowGraphNode ( self . handle , index, node. handle ) }
92+ }
93+
94+ pub fn clear ( & self ) {
95+ unsafe { BNClearFlowGraphNodes ( self . handle ) }
96+ }
97+
4898 pub fn set_option ( & self , option : FlowGraphOption , value : bool ) {
4999 unsafe { BNSetFlowGraphOption ( self . handle , option, value) }
50100 }
51101
52102 pub fn is_option_set ( & self , option : FlowGraphOption ) -> bool {
53103 unsafe { BNIsFlowGraphOptionSet ( self . handle , option) }
54104 }
105+
106+ /// A list of the currently applied [`CoreRenderLayer`]'s
107+ pub fn render_layers ( & self ) -> Array < CoreRenderLayer > {
108+ let mut count: usize = 0 ;
109+ unsafe {
110+ let handles = BNGetFlowGraphRenderLayers ( self . handle , & mut count) ;
111+ Array :: new ( handles, count, ( ) )
112+ }
113+ }
114+
115+ /// Add a Render Layer to be applied to this [`FlowGraph`].
116+ ///
117+ /// NOTE: Layers will be applied in the order in which they are added.
118+ pub fn add_render_layer ( & self , layer : & CoreRenderLayer ) {
119+ unsafe { BNAddFlowGraphRenderLayer ( self . handle , layer. handle . as_ptr ( ) ) } ;
120+ }
121+
122+ /// Remove a Render Layer from being applied to this [`FlowGraph`].
123+ pub fn remove_render_layer ( & self , layer : & CoreRenderLayer ) {
124+ unsafe { BNRemoveFlowGraphRenderLayer ( self . handle , layer. handle . as_ptr ( ) ) } ;
125+ }
55126}
56127
57128unsafe impl RefCountable for FlowGraph {
@@ -88,8 +159,37 @@ impl<'a> FlowGraphNode<'a> {
88159 }
89160 }
90161
91- pub fn new ( graph : & FlowGraph ) -> Self {
92- unsafe { FlowGraphNode :: from_raw ( BNCreateFlowGraphNode ( graph. handle ) ) }
162+ pub ( crate ) unsafe fn ref_from_raw ( raw : * mut BNFlowGraphNode ) -> Ref < Self > {
163+ Ref :: new ( Self {
164+ handle : raw,
165+ _data : PhantomData ,
166+ } )
167+ }
168+
169+ pub fn new ( graph : & FlowGraph ) -> Ref < Self > {
170+ unsafe { FlowGraphNode :: ref_from_raw ( BNCreateFlowGraphNode ( graph. handle ) ) }
171+ }
172+
173+ pub fn basic_block < C : BlockContext > ( & self , context : C ) -> Option < Ref < BasicBlock < C > > > {
174+ let block_ptr = unsafe { BNGetFlowGraphBasicBlock ( self . handle ) } ;
175+ if block_ptr. is_null ( ) {
176+ return None ;
177+ }
178+ Some ( unsafe { BasicBlock :: ref_from_raw ( block_ptr, context) } )
179+ }
180+
181+ pub fn set_basic_block < C : BlockContext > ( & self , block : Option < & BasicBlock < C > > ) {
182+ match block {
183+ Some ( block) => unsafe { BNSetFlowGraphBasicBlock ( self . handle , block. handle ) } ,
184+ None => unsafe { BNSetFlowGraphBasicBlock ( self . handle , std:: ptr:: null_mut ( ) ) } ,
185+ }
186+ }
187+
188+ pub fn lines ( & self ) -> Array < DisassemblyTextLine > {
189+ let mut count = 0 ;
190+ let result = unsafe { BNGetFlowGraphNodeLines ( self . handle , & mut count) } ;
191+ assert ! ( !result. is_null( ) ) ;
192+ unsafe { Array :: new ( result, count, ( ) ) }
93193 }
94194
95195 pub fn set_lines ( & self , lines : impl IntoIterator < Item = DisassemblyTextLine > ) {
@@ -106,6 +206,30 @@ impl<'a> FlowGraphNode<'a> {
106206 }
107207 }
108208
209+ /// Returns the graph position of the node in X, Y form.
210+ pub fn position ( & self ) -> ( i32 , i32 ) {
211+ let pos_x = unsafe { BNGetFlowGraphNodeX ( self . handle ) } ;
212+ let pos_y = unsafe { BNGetFlowGraphNodeY ( self . handle ) } ;
213+ ( pos_x, pos_y)
214+ }
215+
216+ /// Sets the graph position of the node.
217+ pub fn set_position ( & self , x : i32 , y : i32 ) {
218+ unsafe { BNFlowGraphNodeSetX ( self . handle , x) } ;
219+ unsafe { BNFlowGraphNodeSetX ( self . handle , y) } ;
220+ }
221+
222+ pub fn highlight_color ( & self ) -> HighlightColor {
223+ let raw = unsafe { BNGetFlowGraphNodeHighlight ( self . handle ) } ;
224+ HighlightColor :: from ( raw)
225+ }
226+
227+ pub fn set_highlight_color ( & self , highlight : HighlightColor ) {
228+ unsafe { BNSetFlowGraphNodeHighlight ( self . handle , highlight. into ( ) ) } ;
229+ }
230+
231+ // TODO: Add getters and setters for edges
232+
109233 pub fn add_outgoing_edge (
110234 & self ,
111235 type_ : BranchType ,
0 commit comments