threatflux_binary_analysis/analysis/
visualization.rs

1#[cfg(feature = "visualization")]
2#[allow(unused_imports)]
3use dot_generator as _dot_generator;
4
5use crate::types::ControlFlowGraph;
6
7/// Generate a DOT representation of a control flow graph.
8#[cfg(feature = "visualization")]
9pub fn cfg_to_dot(cfg: &ControlFlowGraph) -> String {
10    let mut dot = String::from("digraph cfg {\n");
11    for block in &cfg.basic_blocks {
12        dot.push_str(&format!(
13            "  bb{} [label=\"0x{:x}\"];\n",
14            block.id, block.start_address
15        ));
16    }
17    for block in &cfg.basic_blocks {
18        for succ in &block.successors {
19            dot.push_str(&format!("  bb{} -> bb{};\n", block.id, succ));
20        }
21    }
22    dot.push_str("}\n");
23    dot
24}
25
26#[cfg(test)]
27mod tests {
28    use super::*;
29    use crate::types::{BasicBlock, ComplexityMetrics, Function, FunctionType};
30
31    #[cfg(feature = "visualization")]
32    #[test]
33    fn test_cfg_to_dot() {
34        let cfg = ControlFlowGraph {
35            function: Function {
36                name: "test".into(),
37                start_address: 0,
38                end_address: 10,
39                size: 10,
40                function_type: FunctionType::Normal,
41                calling_convention: None,
42                parameters: vec![],
43                return_type: None,
44            },
45            basic_blocks: vec![
46                BasicBlock {
47                    id: 0,
48                    start_address: 0,
49                    end_address: 5,
50                    instructions: vec![],
51                    successors: vec![1],
52                    predecessors: vec![],
53                    block_type: crate::types::BlockType::Entry,
54                    dominator: None,
55                    dominance_frontier: Vec::new(),
56                },
57                BasicBlock {
58                    id: 1,
59                    start_address: 5,
60                    end_address: 10,
61                    instructions: vec![],
62                    successors: vec![],
63                    predecessors: vec![0],
64                    block_type: crate::types::BlockType::Exit,
65                    dominator: None,
66                    dominance_frontier: Vec::new(),
67                },
68            ],
69            complexity: ComplexityMetrics {
70                cyclomatic_complexity: 1,
71                basic_block_count: 2,
72                edge_count: 1,
73                nesting_depth: 0,
74                loop_count: 0,
75                cognitive_complexity: 1,
76                halstead_metrics: None,
77                maintainability_index: None,
78            },
79            loops: vec![],
80        };
81
82        let dot = cfg_to_dot(&cfg);
83        assert!(dot.contains("bb0 -> bb1"));
84        assert!(dot.contains("bb0"));
85        assert!(dot.contains("bb1"));
86    }
87}