Skip to main content

hawk_core/
graph.rs

1use indexmap::IndexMap;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5// ---------------------------------------------------------------------------
6// NodeKind
7// ---------------------------------------------------------------------------
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
10#[serde(rename_all = "PascalCase")]
11pub enum NodeKind {
12    Lambda,
13    ApiGateway,
14    ApiRoute,
15    EventRule,
16    SqsQueue,
17    SnsTopic,
18    S3Bucket,
19    DynamoStream,
20    StepFunction,
21    LogGroup,
22    EcsService,
23    Ec2Instance,
24    LoadBalancer,
25    Unknown,
26}
27
28impl std::fmt::Display for NodeKind {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        write!(f, "{:?}", self)
31    }
32}
33
34// ---------------------------------------------------------------------------
35// EdgeKind
36// ---------------------------------------------------------------------------
37
38#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
39#[serde(rename_all = "PascalCase")]
40pub enum EdgeKind {
41    Triggers,
42    Invokes,
43    Consumes,
44    Publishes,
45    ReadsFrom,
46    WritesTo,
47}
48
49impl std::fmt::Display for EdgeKind {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        write!(f, "{:?}", self)
52    }
53}
54
55// ---------------------------------------------------------------------------
56// Node
57// ---------------------------------------------------------------------------
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct Node {
61    pub id: String,
62    pub kind: NodeKind,
63    pub name: String,
64    #[serde(skip_serializing_if = "Option::is_none")]
65    pub arn: Option<String>,
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub region: Option<String>,
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub account_id: Option<String>,
70    #[serde(skip_serializing_if = "Option::is_none")]
71    pub tags: Option<IndexMap<String, String>>,
72    #[serde(default = "default_props")]
73    pub props: serde_json::Value,
74}
75
76fn default_props() -> serde_json::Value {
77    serde_json::Value::Object(serde_json::Map::new())
78}
79
80// ---------------------------------------------------------------------------
81// Edge
82// ---------------------------------------------------------------------------
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct Edge {
86    pub from: String,
87    pub to: String,
88    pub kind: EdgeKind,
89    #[serde(default = "default_props")]
90    pub props: serde_json::Value,
91}
92
93// ---------------------------------------------------------------------------
94// GraphStats
95// ---------------------------------------------------------------------------
96
97#[derive(Debug, Clone, Default, Serialize, Deserialize)]
98pub struct GraphStats {
99    pub node_count: usize,
100    pub edge_count: usize,
101    pub nodes_by_kind: IndexMap<String, usize>,
102    pub edges_by_kind: IndexMap<String, usize>,
103    pub top_fan_in: Vec<(String, usize)>,
104    pub top_fan_out: Vec<(String, usize)>,
105}
106
107// ---------------------------------------------------------------------------
108// Graph
109// ---------------------------------------------------------------------------
110
111#[derive(Debug, Clone, Serialize, Deserialize)]
112pub struct Graph {
113    pub generated_at: String,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub profile: Option<String>,
116    pub regions: Vec<String>,
117    pub nodes: Vec<Node>,
118    pub edges: Vec<Edge>,
119    #[serde(default)]
120    pub warnings: Vec<String>,
121    #[serde(default)]
122    pub stats: GraphStats,
123}
124
125impl Default for Graph {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl Graph {
132    pub fn new() -> Self {
133        Self {
134            generated_at: chrono::Utc::now().to_rfc3339(),
135            profile: None,
136            regions: Vec::new(),
137            nodes: Vec::new(),
138            edges: Vec::new(),
139            warnings: Vec::new(),
140            stats: GraphStats::default(),
141        }
142    }
143
144    /// Remove duplicate nodes (by id) and edges (by from+to+kind), then sort
145    /// deterministically.
146    pub fn dedupe_and_sort(&mut self) {
147        // Dedupe nodes by id, keep first occurrence
148        let mut seen = std::collections::HashSet::new();
149        self.nodes.retain(|n| seen.insert(n.id.clone()));
150
151        // Dedupe edges by (from, to, kind)
152        let mut seen_edges = std::collections::HashSet::new();
153        self.edges.retain(|e| {
154            seen_edges.insert((e.from.clone(), e.to.clone(), format!("{:?}", e.kind)))
155        });
156
157        // Sort nodes by kind, then name, then id
158        self.nodes.sort_by(|a, b| {
159            a.kind
160                .cmp(&b.kind)
161                .then_with(|| a.name.cmp(&b.name))
162                .then_with(|| a.id.cmp(&b.id))
163        });
164
165        // Sort edges by kind, then from, then to
166        self.edges.sort_by(|a, b| {
167            a.kind
168                .cmp(&b.kind)
169                .then_with(|| a.from.cmp(&b.from))
170                .then_with(|| a.to.cmp(&b.to))
171        });
172    }
173
174    /// Compute graph statistics.
175    pub fn compute_stats(&mut self) {
176        // Build id -> name lookup for human-readable fan-in/fan-out labels
177        let id_to_name: HashMap<&str, &str> = self
178            .nodes
179            .iter()
180            .map(|n| (n.id.as_str(), n.name.as_str()))
181            .collect();
182
183        let mut nodes_by_kind: IndexMap<String, usize> = IndexMap::new();
184        for node in &self.nodes {
185            *nodes_by_kind.entry(node.kind.to_string()).or_default() += 1;
186        }
187
188        let mut edges_by_kind: IndexMap<String, usize> = IndexMap::new();
189        for edge in &self.edges {
190            *edges_by_kind.entry(edge.kind.to_string()).or_default() += 1;
191        }
192
193        // Fan-in: count incoming edges per node (resolve to name)
194        let mut fan_in: HashMap<String, usize> = HashMap::new();
195        for edge in &self.edges {
196            let label = id_to_name
197                .get(edge.to.as_str())
198                .unwrap_or(&edge.to.as_str())
199                .to_string();
200            *fan_in.entry(label).or_default() += 1;
201        }
202        let mut top_fan_in: Vec<(String, usize)> = fan_in.into_iter().collect();
203        top_fan_in.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
204        top_fan_in.truncate(10);
205
206        // Fan-out: count outgoing edges per node (resolve to name)
207        let mut fan_out: HashMap<String, usize> = HashMap::new();
208        for edge in &self.edges {
209            let label = id_to_name
210                .get(edge.from.as_str())
211                .unwrap_or(&edge.from.as_str())
212                .to_string();
213            *fan_out.entry(label).or_default() += 1;
214        }
215        let mut top_fan_out: Vec<(String, usize)> = fan_out.into_iter().collect();
216        top_fan_out.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
217        top_fan_out.truncate(10);
218
219        self.stats = GraphStats {
220            node_count: self.nodes.len(),
221            edge_count: self.edges.len(),
222            nodes_by_kind,
223            edges_by_kind,
224            top_fan_in,
225            top_fan_out,
226        };
227    }
228
229    /// Merge another graph's output into this graph.
230    pub fn merge(&mut self, output: DiscoveryOutput) {
231        self.nodes.extend(output.nodes);
232        self.edges.extend(output.edges);
233        self.warnings.extend(output.warnings);
234    }
235}
236
237// ---------------------------------------------------------------------------
238// DiscoveryOutput — returned by each discovery module
239// ---------------------------------------------------------------------------
240
241#[derive(Debug, Clone, Default)]
242pub struct DiscoveryOutput {
243    pub nodes: Vec<Node>,
244    pub edges: Vec<Edge>,
245    pub warnings: Vec<String>,
246}
247
248// ---------------------------------------------------------------------------
249// Diff support
250// ---------------------------------------------------------------------------
251
252#[derive(Debug, Clone, Serialize, Deserialize)]
253pub struct GraphDiff {
254    pub added_nodes: Vec<String>,
255    pub removed_nodes: Vec<String>,
256    pub added_edges: Vec<String>,
257    pub removed_edges: Vec<String>,
258}
259
260impl GraphDiff {
261    pub fn compute(old: &Graph, new: &Graph) -> Self {
262        let old_node_ids: std::collections::HashSet<&str> =
263            old.nodes.iter().map(|n| n.id.as_str()).collect();
264        let new_node_ids: std::collections::HashSet<&str> =
265            new.nodes.iter().map(|n| n.id.as_str()).collect();
266
267        let added_nodes: Vec<String> = new_node_ids
268            .difference(&old_node_ids)
269            .map(|s| s.to_string())
270            .collect();
271        let removed_nodes: Vec<String> = old_node_ids
272            .difference(&new_node_ids)
273            .map(|s| s.to_string())
274            .collect();
275
276        let edge_key = |e: &Edge| format!("{} --{:?}--> {}", e.from, e.kind, e.to);
277        let old_edge_keys: std::collections::HashSet<String> =
278            old.edges.iter().map(edge_key).collect();
279        let new_edge_keys: std::collections::HashSet<String> =
280            new.edges.iter().map(edge_key).collect();
281
282        let added_edges: Vec<String> = new_edge_keys
283            .difference(&old_edge_keys)
284            .cloned()
285            .collect();
286        let removed_edges: Vec<String> = old_edge_keys
287            .difference(&new_edge_keys)
288            .cloned()
289            .collect();
290
291        GraphDiff {
292            added_nodes,
293            removed_nodes,
294            added_edges,
295            removed_edges,
296        }
297    }
298}
299
300// ---------------------------------------------------------------------------
301// Tests
302// ---------------------------------------------------------------------------
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn test_dedupe_and_sort() {
310        let mut g = Graph::new();
311        g.nodes.push(Node {
312            id: "arn:aws:lambda:us-east-1:123:function:alpha".into(),
313            kind: NodeKind::Lambda,
314            name: "alpha".into(),
315            arn: Some("arn:aws:lambda:us-east-1:123:function:alpha".into()),
316            region: Some("us-east-1".into()),
317            account_id: None,
318            tags: None,
319            props: serde_json::json!({}),
320        });
321        // Duplicate
322        g.nodes.push(Node {
323            id: "arn:aws:lambda:us-east-1:123:function:alpha".into(),
324            kind: NodeKind::Lambda,
325            name: "alpha".into(),
326            arn: Some("arn:aws:lambda:us-east-1:123:function:alpha".into()),
327            region: Some("us-east-1".into()),
328            account_id: None,
329            tags: None,
330            props: serde_json::json!({}),
331        });
332        g.nodes.push(Node {
333            id: "arn:aws:lambda:us-east-1:123:function:beta".into(),
334            kind: NodeKind::Lambda,
335            name: "beta".into(),
336            arn: None,
337            region: None,
338            account_id: None,
339            tags: None,
340            props: serde_json::json!({}),
341        });
342
343        g.edges.push(Edge {
344            from: "a".into(),
345            to: "b".into(),
346            kind: EdgeKind::Triggers,
347            props: serde_json::json!({}),
348        });
349        // Duplicate edge
350        g.edges.push(Edge {
351            from: "a".into(),
352            to: "b".into(),
353            kind: EdgeKind::Triggers,
354            props: serde_json::json!({}),
355        });
356
357        g.dedupe_and_sort();
358        assert_eq!(g.nodes.len(), 2);
359        assert_eq!(g.edges.len(), 1);
360        // alpha before beta
361        assert_eq!(g.nodes[0].name, "alpha");
362        assert_eq!(g.nodes[1].name, "beta");
363    }
364
365    #[test]
366    fn test_compute_stats() {
367        let mut g = Graph::new();
368        g.nodes.push(Node {
369            id: "l1".into(),
370            kind: NodeKind::Lambda,
371            name: "fn1".into(),
372            arn: None,
373            region: None,
374            account_id: None,
375            tags: None,
376            props: serde_json::json!({}),
377        });
378        g.nodes.push(Node {
379            id: "l2".into(),
380            kind: NodeKind::Lambda,
381            name: "fn2".into(),
382            arn: None,
383            region: None,
384            account_id: None,
385            tags: None,
386            props: serde_json::json!({}),
387        });
388        g.nodes.push(Node {
389            id: "q1".into(),
390            kind: NodeKind::SqsQueue,
391            name: "queue1".into(),
392            arn: None,
393            region: None,
394            account_id: None,
395            tags: None,
396            props: serde_json::json!({}),
397        });
398        g.edges.push(Edge {
399            from: "q1".into(),
400            to: "l1".into(),
401            kind: EdgeKind::Triggers,
402            props: serde_json::json!({}),
403        });
404        g.edges.push(Edge {
405            from: "q1".into(),
406            to: "l2".into(),
407            kind: EdgeKind::Triggers,
408            props: serde_json::json!({}),
409        });
410        g.compute_stats();
411        assert_eq!(g.stats.node_count, 3);
412        assert_eq!(g.stats.edge_count, 2);
413        assert_eq!(g.stats.nodes_by_kind["Lambda"], 2);
414        assert_eq!(g.stats.top_fan_out[0], ("queue1".to_string(), 2));
415    }
416
417    #[test]
418    fn test_diff() {
419        let mut old = Graph::new();
420        old.nodes.push(Node {
421            id: "a".into(),
422            kind: NodeKind::Lambda,
423            name: "a".into(),
424            arn: None,
425            region: None,
426            account_id: None,
427            tags: None,
428            props: serde_json::json!({}),
429        });
430
431        let mut new = Graph::new();
432        new.nodes.push(Node {
433            id: "b".into(),
434            kind: NodeKind::Lambda,
435            name: "b".into(),
436            arn: None,
437            region: None,
438            account_id: None,
439            tags: None,
440            props: serde_json::json!({}),
441        });
442
443        let diff = GraphDiff::compute(&old, &new);
444        assert_eq!(diff.added_nodes, vec!["b".to_string()]);
445        assert_eq!(diff.removed_nodes, vec!["a".to_string()]);
446    }
447
448    #[test]
449    fn test_graph_serialization_roundtrip() {
450        let mut g = Graph::new();
451        g.nodes.push(Node {
452            id: "test".into(),
453            kind: NodeKind::Lambda,
454            name: "test-fn".into(),
455            arn: Some("arn:aws:lambda:us-east-1:123:function:test-fn".into()),
456            region: Some("us-east-1".into()),
457            account_id: Some("123".into()),
458            tags: None,
459            props: serde_json::json!({"runtime": "nodejs18.x"}),
460        });
461        g.compute_stats();
462
463        let json = serde_json::to_string_pretty(&g).unwrap();
464        let parsed: Graph = serde_json::from_str(&json).unwrap();
465        assert_eq!(parsed.nodes.len(), 1);
466        assert_eq!(parsed.nodes[0].name, "test-fn");
467    }
468}