1use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ThreatIntelConfig {
9 pub enabled: bool,
10 pub sources: HashMap<String, SourceConfig>,
11 pub sync_interval_hours: u32,
12 pub cache_enabled: bool,
13 pub cache_ttl_hours: u32,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct SourceConfig {
19 pub id: String,
20 pub name: String,
21 pub source_type: SourceType,
22 pub enabled: bool,
23 pub api_url: Option<String>,
24 pub api_key: Option<String>,
25 pub auth_type: AuthType,
26 pub update_frequency: UpdateFrequency,
27 pub priority: u32,
28 pub capabilities: Vec<SourceCapability>,
29 pub timeout_secs: u64,
30 pub retry_count: u32,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
35#[serde(rename_all = "snake_case")]
36pub enum SourceType {
37 MitreAttack,
38 Cve,
39 Osint,
40 Commercial,
41 Custom,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
46#[serde(rename_all = "snake_case")]
47pub enum AuthType {
48 None,
49 ApiKey,
50 Bearer,
51 Basic,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
56#[serde(rename_all = "snake_case")]
57pub enum UpdateFrequency {
58 Realtime,
59 Hourly,
60 Daily,
61 Weekly,
62 Manual,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
67#[serde(rename_all = "snake_case")]
68pub enum SourceCapability {
69 Vulnerabilities,
70 Exploits,
71 ThreatActors,
72 Ioc,
73 Advisories,
74 Tactics,
75 Techniques,
76 Malware,
77 Patches,
78}
79
80impl Default for ThreatIntelConfig {
81 fn default() -> Self {
82 let mut sources = HashMap::new();
83
84 sources.insert(
86 "mitre_attack".to_string(),
87 SourceConfig {
88 id: "mitre_attack".to_string(),
89 name: "MITRE ATT&CK".to_string(),
90 source_type: SourceType::MitreAttack,
91 enabled: true,
92 api_url: Some(
93 "https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json"
94 .to_string(),
95 ),
96 api_key: None,
97 auth_type: AuthType::None,
98 update_frequency: UpdateFrequency::Daily,
99 priority: 10,
100 capabilities: vec![
101 SourceCapability::ThreatActors,
102 SourceCapability::Tactics,
103 SourceCapability::Techniques,
104 ],
105 timeout_secs: 30,
106 retry_count: 3,
107 },
108 );
109
110 sources.insert(
112 "cve_database".to_string(),
113 SourceConfig {
114 id: "cve_database".to_string(),
115 name: "CVE Database".to_string(),
116 source_type: SourceType::Cve,
117 enabled: true,
118 api_url: Some("https://services.nvd.nist.gov/rest/json/cves/2.0".to_string()),
119 api_key: None,
120 auth_type: AuthType::None,
121 update_frequency: UpdateFrequency::Realtime,
122 priority: 9,
123 capabilities: vec![
124 SourceCapability::Vulnerabilities,
125 SourceCapability::Exploits,
126 SourceCapability::Patches,
127 ],
128 timeout_secs: 60,
129 retry_count: 3,
130 },
131 );
132
133 sources.insert(
135 "abuse_ch".to_string(),
136 SourceConfig {
137 id: "abuse_ch".to_string(),
138 name: "Abuse.ch".to_string(),
139 source_type: SourceType::Osint,
140 enabled: true,
141 api_url: Some("https://urlhaus-api.abuse.ch/v1/urls/recent/".to_string()),
142 api_key: None,
143 auth_type: AuthType::None,
144 update_frequency: UpdateFrequency::Hourly,
145 priority: 7,
146 capabilities: vec![SourceCapability::Ioc, SourceCapability::Malware],
147 timeout_secs: 30,
148 retry_count: 2,
149 },
150 );
151
152 Self {
153 enabled: true,
154 sources,
155 sync_interval_hours: 24,
156 cache_enabled: true,
157 cache_ttl_hours: 6,
158 }
159 }
160}
161
162impl ThreatIntelConfig {
163 pub fn new() -> Self {
165 Self {
166 enabled: true,
167 sources: HashMap::new(),
168 sync_interval_hours: 24,
169 cache_enabled: true,
170 cache_ttl_hours: 6,
171 }
172 }
173
174 pub fn get_enabled_sources(&self) -> Vec<&SourceConfig> {
176 let mut sources: Vec<&SourceConfig> =
177 self.sources.values().filter(|s| s.enabled).collect();
178
179 sources.sort_by(|a, b| b.priority.cmp(&a.priority));
180 sources
181 }
182
183 pub fn get_sources_by_capability(&self, capability: SourceCapability) -> Vec<&SourceConfig> {
185 self.sources
186 .values()
187 .filter(|s| s.enabled && s.capabilities.contains(&capability))
188 .collect()
189 }
190
191 pub fn add_source(&mut self, source: SourceConfig) {
193 self.sources.insert(source.id.clone(), source);
194 }
195
196 pub fn remove_source(&mut self, id: &str) -> Option<SourceConfig> {
198 self.sources.remove(id)
199 }
200
201 pub fn set_source_enabled(&mut self, id: &str, enabled: bool) -> bool {
203 if let Some(source) = self.sources.get_mut(id) {
204 source.enabled = enabled;
205 true
206 } else {
207 false
208 }
209 }
210
211 pub fn get_source(&self, id: &str) -> Option<&SourceConfig> {
213 self.sources.get(id)
214 }
215
216 pub fn get_source_mut(&mut self, id: &str) -> Option<&mut SourceConfig> {
218 self.sources.get_mut(id)
219 }
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225
226 #[test]
227 fn test_default_config() {
228 let config = ThreatIntelConfig::default();
229 assert!(config.enabled);
230 assert!(config.sources.len() >= 3);
231 assert!(config.sources.contains_key("mitre_attack"));
232 assert!(config.sources.contains_key("cve_database"));
233 assert!(config.sources.contains_key("abuse_ch"));
234 }
235
236 #[test]
237 fn test_get_enabled_sources() {
238 let config = ThreatIntelConfig::default();
239 let sources = config.get_enabled_sources();
240
241 assert_eq!(sources.len(), 3);
243
244 for i in 0..sources.len() - 1 {
246 assert!(sources[i].priority >= sources[i + 1].priority);
247 }
248 }
249
250 #[test]
251 fn test_get_sources_by_capability() {
252 let config = ThreatIntelConfig::default();
253
254 let vuln_sources = config.get_sources_by_capability(SourceCapability::Vulnerabilities);
255 assert!(!vuln_sources.is_empty());
256
257 for source in vuln_sources {
258 assert!(source.capabilities.contains(&SourceCapability::Vulnerabilities));
259 }
260
261 let ioc_sources = config.get_sources_by_capability(SourceCapability::Ioc);
262 assert!(!ioc_sources.is_empty());
263
264 for source in ioc_sources {
265 assert!(source.capabilities.contains(&SourceCapability::Ioc));
266 }
267 }
268
269 #[test]
270 fn test_add_remove_source() {
271 let mut config = ThreatIntelConfig::new();
272
273 let custom_source = SourceConfig {
274 id: "custom".to_string(),
275 name: "Custom Source".to_string(),
276 source_type: SourceType::Custom,
277 enabled: true,
278 api_url: Some("https://example.com/api".to_string()),
279 api_key: Some("test-key".to_string()),
280 auth_type: AuthType::ApiKey,
281 update_frequency: UpdateFrequency::Daily,
282 priority: 5,
283 capabilities: vec![SourceCapability::Ioc],
284 timeout_secs: 30,
285 retry_count: 3,
286 };
287
288 config.add_source(custom_source.clone());
289 assert!(config.sources.contains_key("custom"));
290 assert_eq!(config.sources.get("custom").unwrap().name, "Custom Source");
291
292 let removed = config.remove_source("custom");
293 assert!(removed.is_some());
294 assert!(!config.sources.contains_key("custom"));
295 }
296
297 #[test]
298 fn test_enable_disable_source() {
299 let mut config = ThreatIntelConfig::default();
300
301 let success = config.set_source_enabled("mitre_attack", false);
302 assert!(success);
303 assert!(!config.sources.get("mitre_attack").unwrap().enabled);
304
305 let success = config.set_source_enabled("mitre_attack", true);
306 assert!(success);
307 assert!(config.sources.get("mitre_attack").unwrap().enabled);
308
309 let failure = config.set_source_enabled("nonexistent", true);
310 assert!(!failure);
311 }
312
313 #[test]
314 fn test_get_source() {
315 let config = ThreatIntelConfig::default();
316
317 let source = config.get_source("mitre_attack");
318 assert!(source.is_some());
319 assert_eq!(source.unwrap().name, "MITRE ATT&CK");
320
321 let missing = config.get_source("nonexistent");
322 assert!(missing.is_none());
323 }
324}
325