Skip to main content

threatflux_binary_analysis/analysis/
visualization.rs

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