attack/store/
mod.rs

1use std::collections::HashMap;
2use std::path::Path;
3use serde::Deserialize;
4use crate::domain::{*, software};
5use crate::error::Result;
6use stix_rs::{Relationship, StixObject};
7
8#[derive(Debug, Default)]
9pub struct AttackStore {
10    tactics: HashMap<String, Tactic>,
11    techniques: HashMap<String, Technique>,
12    groups: HashMap<String, Group>,
13    software: HashMap<String, Software>,
14    mitigations: HashMap<String, Mitigation>,
15    data_sources: HashMap<String, DataSource>,
16    data_components: HashMap<String, DataComponent>,
17    campaigns: HashMap<String, Campaign>,
18    matrices: HashMap<String, Matrix>,
19    analytics: HashMap<String, Analytic>,
20    detection_strategies: HashMap<String, DetectionStrategy>,
21
22    // Secondary Indices
23    tcode_index: HashMap<String, String>, // T-Code -> ID
24    technique_to_group: HashMap<String, Vec<String>>,
25    software_to_group: HashMap<String, Vec<String>>,
26    technique_to_mitigation: HashMap<String, Vec<String>>,
27    technique_to_datasource: HashMap<String, Vec<String>>,
28    technique_to_datacomponent: HashMap<String, Vec<String>>,
29    subtechnique_to_technique: HashMap<String, String>,
30    campaign_to_group: HashMap<String, Vec<String>>,
31    campaign_to_technique: HashMap<String, Vec<String>>,
32    campaign_to_software: HashMap<String, Vec<String>>,
33    datacomponent_to_datasource: HashMap<String, String>,
34
35    // Replacement Index (source_ref -> target_ref where type is revoked-by)
36    replacements: HashMap<String, String>,
37}
38
39#[derive(Deserialize)]
40struct StixBundle {
41    objects: Vec<serde_json::Value>,
42}
43
44impl AttackStore {
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
50        let file = std::fs::File::open(path)?;
51        let bundle: StixBundle = serde_json::from_reader(file)?;
52        Ok(Self::from_bundle(bundle))
53    }
54
55    pub fn from_url(url: &str) -> Result<Self> {
56        let resp = reqwest::blocking::get(url)?;
57        let bundle: StixBundle = resp.json()?;
58        Ok(Self::from_bundle(bundle))
59    }
60
61    fn from_bundle(bundle: StixBundle) -> Self {
62        let mut store = Self::new();
63        let mut relationships = Vec::new();
64
65        for obj in bundle.objects {
66            let type_ = obj.get("type").and_then(|v| v.as_str()).unwrap_or("");
67            match type_ {
68                "x-mitre-tactic" => {
69                    if let Ok(tactic) = serde_json::from_value::<Tactic>(obj) {
70                        store.tactics.insert(tactic.id().to_string(), tactic);
71                    }
72                }
73                "attack-pattern" => {
74                    if let Ok(technique) = serde_json::from_value::<Technique>(obj) {
75                        if let Some(tcode) = technique.tcode() {
76                            store.tcode_index.insert(tcode.to_string(), technique.id().to_string());
77                        }
78                        store.techniques.insert(technique.id().to_string(), technique);
79                    }
80                }
81                "intrusion-set" => {
82                    if let Ok(group) = serde_json::from_value::<Group>(obj) {
83                        store.groups.insert(group.id().to_string(), group);
84                    }
85                }
86                "malware" => {
87                    if let Ok(malware) = serde_json::from_value::<software::Malware>(obj) {
88                        let id = malware.id().to_string();
89                        store.software.insert(id, Software::Malware(malware));
90                    }
91                }
92                "tool" => {
93                    if let Ok(tool) = serde_json::from_value::<software::Tool>(obj) {
94                        let id = tool.id().to_string();
95                        store.software.insert(id, Software::Tool(tool));
96                    }
97                }
98                "course-of-action" => {
99                    if let Ok(mitigation) = serde_json::from_value::<Mitigation>(obj) {
100                        store.mitigations.insert(mitigation.id().to_string(), mitigation);
101                    }
102                }
103                "x-mitre-data-source" => {
104                    if let Ok(ds) = serde_json::from_value::<DataSource>(obj) {
105                        store.data_sources.insert(ds.id().to_string(), ds);
106                    }
107                }
108                "x-mitre-data-component" => {
109                    if let Ok(dc) = serde_json::from_value::<DataComponent>(obj.clone()) {
110                        // Store data component to data source mapping
111                        if let Some(ds_ref) = &dc.data_source_ref {
112                            store.datacomponent_to_datasource.insert(dc.id().to_string(), ds_ref.clone());
113                        }
114                        store.data_components.insert(dc.id().to_string(), dc);
115                    }
116                }
117                "campaign" => {
118                    if let Ok(campaign) = serde_json::from_value::<Campaign>(obj) {
119                        store.campaigns.insert(campaign.id().to_string(), campaign);
120                    }
121                }
122                "x-mitre-matrix" => {
123                    if let Ok(matrix) = serde_json::from_value::<Matrix>(obj) {
124                        store.matrices.insert(matrix.id().to_string(), matrix);
125                    }
126                }
127                "x-mitre-analytic" => {
128                    if let Ok(analytic) = serde_json::from_value::<Analytic>(obj) {
129                        store.analytics.insert(analytic.id().to_string(), analytic);
130                    }
131                }
132                "x-mitre-detection-strategy" => {
133                    if let Ok(strategy) = serde_json::from_value::<DetectionStrategy>(obj) {
134                        store.detection_strategies.insert(strategy.id().to_string(), strategy);
135                    }
136                }
137                "relationship" => {
138                    if let Ok(rel) = serde_json::from_value::<Relationship>(obj) {
139                        relationships.push(rel);
140                    }
141                }
142                _ => {}
143            }
144        }
145
146        // Process relationships and build parent-child hierarchy
147        for tech in store.techniques.values() {
148            if tech.is_subtechnique {
149                // Extract parent from T-Code (e.g., T1566.001 -> T1566)
150                if let Some(tcode) = tech.tcode()
151                    && let Some(parent_tcode) = tcode.split('.').next()
152                    && let Some(parent_id) = store.tcode_index.get(parent_tcode)
153                {
154                    store.subtechnique_to_technique.insert(tech.id().to_string(), parent_id.clone());
155                }
156            }
157        }
158
159        // Process relationships
160        for rel in relationships {
161            match rel.relationship_type.as_str() {
162                "uses" => {
163                    // Group/Campaign -> Technique
164                    if rel.target_ref.starts_with("attack-pattern--") {
165                        if rel.source_ref.starts_with("intrusion-set--") {
166                            store.technique_to_group.entry(rel.target_ref.clone()).or_default().push(rel.source_ref.clone());
167                        } else if rel.source_ref.starts_with("campaign--") {
168                            store.campaign_to_technique.entry(rel.source_ref.clone()).or_default().push(rel.target_ref.clone());
169                        }
170                    }
171                    // Group/Campaign -> Software
172                    else if rel.target_ref.starts_with("malware--") || rel.target_ref.starts_with("tool--") {
173                        if rel.source_ref.starts_with("intrusion-set--") {
174                            store.software_to_group.entry(rel.target_ref.clone()).or_default().push(rel.source_ref.clone());
175                        } else if rel.source_ref.starts_with("campaign--") {
176                            store.campaign_to_software.entry(rel.source_ref.clone()).or_default().push(rel.target_ref.clone());
177                        }
178                    }
179                }
180                "mitigates" => {
181                    if rel.target_ref.starts_with("attack-pattern--") && rel.source_ref.starts_with("course-of-action--") {
182                        store.technique_to_mitigation.entry(rel.target_ref.clone()).or_default().push(rel.source_ref.clone());
183                    }
184                }
185                "detects" => {
186                    if rel.target_ref.starts_with("attack-pattern--") {
187                        if rel.source_ref.starts_with("x-mitre-data-source--") {
188                            store.technique_to_datasource.entry(rel.target_ref.clone()).or_default().push(rel.source_ref.clone());
189                        } else if rel.source_ref.starts_with("x-mitre-data-component--") {
190                            store.technique_to_datacomponent.entry(rel.target_ref.clone()).or_default().push(rel.source_ref.clone());
191                        }
192                    }
193                }
194                "attributed-to" => {
195                    if rel.source_ref.starts_with("campaign--") && rel.target_ref.starts_with("intrusion-set--") {
196                        store.campaign_to_group.entry(rel.source_ref.clone()).or_default().push(rel.target_ref.clone());
197                    }
198                }
199                "subtechnique-of" => {
200                    if rel.source_ref.starts_with("attack-pattern--") && rel.target_ref.starts_with("attack-pattern--") {
201                        store.subtechnique_to_technique.insert(rel.source_ref.clone(), rel.target_ref.clone());
202                    }
203                }
204                "revoked-by" => {
205                    store.replacements.insert(rel.source_ref.clone(), rel.target_ref.clone());
206                }
207                _ => {}
208            }
209        }
210
211        store
212    }
213
214    pub fn get_tactic(&self, id: &str) -> Option<&Tactic> { self.tactics.get(id) }
215    pub fn get_technique(&self, id: &str) -> Option<&Technique> { self.techniques.get(id) }
216    pub fn get_group(&self, id: &str) -> Option<&Group> { self.groups.get(id) }
217    pub fn get_software(&self, id: &str) -> Option<&Software> { self.software.get(id) }
218    pub fn get_mitigation(&self, id: &str) -> Option<&Mitigation> { self.mitigations.get(id) }
219    pub fn get_data_source(&self, id: &str) -> Option<&DataSource> { self.data_sources.get(id) }
220    pub fn get_data_component(&self, id: &str) -> Option<&DataComponent> { self.data_components.get(id) }
221    pub fn get_campaign(&self, id: &str) -> Option<&Campaign> { self.campaigns.get(id) }
222    pub fn get_matrix(&self, id: &str) -> Option<&Matrix> { self.matrices.get(id) }
223    pub fn get_analytic(&self, id: &str) -> Option<&Analytic> { self.analytics.get(id) }
224    pub fn get_detection_strategy(&self, id: &str) -> Option<&DetectionStrategy> { self.detection_strategies.get(id) }
225
226    pub fn get_replacement(&self, id: &str) -> Option<&str> {
227        self.replacements.get(id).map(|s| s.as_str())
228    }
229
230    pub fn get_technique_by_tcode(&self, tcode: &str) -> Option<&Technique> {
231        let id = self.tcode_index.get(tcode)?;
232        self.techniques.get(id)
233    }
234
235    pub fn get_groups_using_technique(&self, technique_id: &str) -> Vec<&Group> {
236        self.technique_to_group.get(technique_id)
237            .map(|ids| ids.iter().filter_map(|id| self.get_group(id)).collect())
238            .unwrap_or_default()
239    }
240
241    pub fn get_mitigations_for_technique(&self, technique_id: &str) -> Vec<&Mitigation> {
242        self.technique_to_mitigation.get(technique_id)
243            .map(|ids| ids.iter().filter_map(|id| self.get_mitigation(id)).collect())
244            .unwrap_or_default()
245    }
246
247    pub fn get_datasources_for_technique(&self, technique_id: &str) -> Vec<&DataSource> {
248        self.technique_to_datasource.get(technique_id)
249            .map(|ids| ids.iter().filter_map(|id| self.get_data_source(id)).collect())
250            .unwrap_or_default()
251    }
252
253    pub fn get_datacomponents_for_technique(&self, technique_id: &str) -> Vec<&DataComponent> {
254        self.technique_to_datacomponent.get(technique_id)
255            .map(|ids| ids.iter().filter_map(|id| self.get_data_component(id)).collect())
256            .unwrap_or_default()
257    }
258
259    pub fn get_datasource_for_component(&self, component_id: &str) -> Option<&DataSource> {
260        let ds_id = self.datacomponent_to_datasource.get(component_id)?;
261        self.data_sources.get(ds_id)
262    }
263
264    pub fn get_components_for_datasource(&self, datasource_id: &str) -> Vec<&DataComponent> {
265        self.datacomponent_to_datasource.iter()
266            .filter(|(_, ds_id)| ds_id.as_str() == datasource_id)
267            .filter_map(|(dc_id, _)| self.data_components.get(dc_id))
268            .collect()
269    }
270
271    pub fn get_parent_technique(&self, subtechnique_id: &str) -> Option<&Technique> {
272        let parent_id = self.subtechnique_to_technique.get(subtechnique_id)?;
273        self.techniques.get(parent_id)
274    }
275
276    pub fn get_subtechniques(&self, technique_id: &str) -> Vec<&Technique> {
277        self.subtechnique_to_technique.iter()
278            .filter(|(_, parent_id)| parent_id.as_str() == technique_id)
279            .filter_map(|(sub_id, _)| self.techniques.get(sub_id))
280            .collect()
281    }
282
283    pub fn get_groups_for_campaign(&self, campaign_id: &str) -> Vec<&Group> {
284        self.campaign_to_group.get(campaign_id)
285            .map(|ids| ids.iter().filter_map(|id| self.get_group(id)).collect())
286            .unwrap_or_default()
287    }
288
289    pub fn get_techniques_for_campaign(&self, campaign_id: &str) -> Vec<&Technique> {
290        self.campaign_to_technique.get(campaign_id)
291            .map(|ids| ids.iter().filter_map(|id| self.get_technique(id)).collect())
292            .unwrap_or_default()
293    }
294
295    pub fn get_software_for_campaign(&self, campaign_id: &str) -> Vec<&Software> {
296        self.campaign_to_software.get(campaign_id)
297            .map(|ids| ids.iter().filter_map(|id| self.get_software(id)).collect())
298            .unwrap_or_default()
299    }
300
301    pub fn find_technique_by_name(&self, name: &str) -> Option<&Technique> {
302        let name_lower = name.to_lowercase();
303
304        // First try exact match
305        if let Some(t) = self.techniques.values()
306            .find(|t| t.name.to_lowercase() == name_lower) {
307            return Some(t);
308        }
309
310        // Fall back to partial match
311        self.techniques.values()
312            .find(|t| t.name.to_lowercase().contains(&name_lower))
313    }
314
315    pub fn tactics(&self) -> impl Iterator<Item = &Tactic> { self.tactics.values() }
316    pub fn techniques(&self) -> impl Iterator<Item = &Technique> { self.techniques.values() }
317    pub fn groups(&self) -> impl Iterator<Item = &Group> { self.groups.values() }
318    pub fn software(&self) -> impl Iterator<Item = &Software> { self.software.values() }
319    pub fn mitigations(&self) -> impl Iterator<Item = &Mitigation> { self.mitigations.values() }
320    pub fn data_sources(&self) -> impl Iterator<Item = &DataSource> { self.data_sources.values() }
321    pub fn data_components(&self) -> impl Iterator<Item = &DataComponent> { self.data_components.values() }
322    pub fn campaigns(&self) -> impl Iterator<Item = &Campaign> { self.campaigns.values() }
323    pub fn matrices(&self) -> impl Iterator<Item = &Matrix> { self.matrices.values() }
324    pub fn analytics(&self) -> impl Iterator<Item = &Analytic> { self.analytics.values() }
325    pub fn detection_strategies(&self) -> impl Iterator<Item = &DetectionStrategy> { self.detection_strategies.values() }
326}
attack/store/mod.rs

attack/store/
mod.rs
