rust_threat_detector/
threat_intelligence.rs

1//! # Threat Intelligence Module
2//!
3//! Indicators of Compromise (IOC) matching and threat intelligence integration
4//! for identifying known malicious actors, IPs, domains, and file hashes.
5
6use crate::{LogEntry, ThreatAlert, ThreatCategory, ThreatSeverity};
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::collections::{HashMap, HashSet};
10
11/// Type of indicator of compromise
12#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
13pub enum IOCType {
14    IPAddress,
15    Domain,
16    FileHash,
17    URL,
18    Email,
19    UserAgent,
20}
21
22/// Indicator of Compromise
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct IOC {
25    pub ioc_type: IOCType,
26    pub value: String,
27    pub severity: ThreatSeverity,
28    pub description: String,
29    pub source: String, // Threat feed source
30    pub first_seen: DateTime<Utc>,
31    pub last_seen: DateTime<Utc>,
32    pub confidence: f64, // 0.0 to 1.0
33}
34
35impl IOC {
36    pub fn new(
37        ioc_type: IOCType,
38        value: String,
39        severity: ThreatSeverity,
40        description: String,
41        source: String,
42    ) -> Self {
43        Self {
44            ioc_type,
45            value,
46            severity,
47            description,
48            source,
49            first_seen: Utc::now(),
50            last_seen: Utc::now(),
51            confidence: 0.8, // Default confidence
52        }
53    }
54}
55
56/// Threat intelligence database
57pub struct ThreatIntelligence {
58    iocs: HashMap<IOCType, HashMap<String, IOC>>,
59    malicious_ips: HashSet<String>,
60    malicious_domains: HashSet<String>,
61    malicious_hashes: HashSet<String>,
62    threat_actors: HashMap<String, ThreatActor>,
63    matches_count: usize,
64}
65
66/// Known threat actor/group
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct ThreatActor {
69    pub name: String,
70    pub aliases: Vec<String>,
71    pub first_seen: DateTime<Utc>,
72    pub techniques: Vec<String>, // MITRE ATT&CK techniques
73    pub targeted_sectors: Vec<String>,
74    pub associated_iocs: Vec<String>,
75}
76
77impl ThreatIntelligence {
78    /// Create new threat intelligence database
79    pub fn new() -> Self {
80        let mut intel = Self {
81            iocs: HashMap::new(),
82            malicious_ips: HashSet::new(),
83            malicious_domains: HashSet::new(),
84            malicious_hashes: HashSet::new(),
85            threat_actors: HashMap::new(),
86            matches_count: 0,
87        };
88
89        // Initialize IOC type maps
90        intel.iocs.insert(IOCType::IPAddress, HashMap::new());
91        intel.iocs.insert(IOCType::Domain, HashMap::new());
92        intel.iocs.insert(IOCType::FileHash, HashMap::new());
93        intel.iocs.insert(IOCType::URL, HashMap::new());
94        intel.iocs.insert(IOCType::Email, HashMap::new());
95        intel.iocs.insert(IOCType::UserAgent, HashMap::new());
96
97        // Load default threat intelligence
98        intel.load_default_iocs();
99
100        intel
101    }
102
103    /// Load default IOCs (example threat intelligence)
104    fn load_default_iocs(&mut self) {
105        // Add known malicious IPs (examples - these would come from threat feeds)
106        self.add_ioc(IOC::new(
107            IOCType::IPAddress,
108            "185.220.101.1".to_string(),
109            ThreatSeverity::High,
110            "Tor exit node - potential anonymization".to_string(),
111            "TorProject".to_string(),
112        ));
113
114        self.add_ioc(IOC::new(
115            IOCType::IPAddress,
116            "45.142.214.0".to_string(),
117            ThreatSeverity::Critical,
118            "Known C2 server IP".to_string(),
119            "ThreatFeed".to_string(),
120        ));
121
122        // Add malicious domains
123        self.add_ioc(IOC::new(
124            IOCType::Domain,
125            "malicious-example.com".to_string(),
126            ThreatSeverity::Critical,
127            "Known phishing domain".to_string(),
128            "PhishTank".to_string(),
129        ));
130
131        // Add malicious file hashes (SHA-256)
132        self.add_ioc(IOC::new(
133            IOCType::FileHash,
134            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855".to_string(),
135            ThreatSeverity::Critical,
136            "Known malware hash".to_string(),
137            "VirusTotal".to_string(),
138        ));
139
140        // Add suspicious user agents
141        self.add_ioc(IOC::new(
142            IOCType::UserAgent,
143            "sqlmap".to_string(),
144            ThreatSeverity::High,
145            "Automated SQL injection tool".to_string(),
146            "SecurityTools".to_string(),
147        ));
148
149        self.add_ioc(IOC::new(
150            IOCType::UserAgent,
151            "nikto".to_string(),
152            ThreatSeverity::High,
153            "Automated vulnerability scanner".to_string(),
154            "SecurityTools".to_string(),
155        ));
156    }
157
158    /// Add an IOC to the database
159    pub fn add_ioc(&mut self, ioc: IOC) {
160        let value = ioc.value.clone();
161
162        // Add to appropriate set for fast lookup
163        match ioc.ioc_type {
164            IOCType::IPAddress => {
165                self.malicious_ips.insert(value.clone());
166            }
167            IOCType::Domain => {
168                self.malicious_domains.insert(value.clone());
169            }
170            IOCType::FileHash => {
171                self.malicious_hashes.insert(value.clone());
172            }
173            _ => {}
174        }
175
176        // Add to main IOC database
177        if let Some(type_map) = self.iocs.get_mut(&ioc.ioc_type) {
178            type_map.insert(value, ioc);
179        }
180    }
181
182    /// Check log against threat intelligence
183    pub fn check_log(&mut self, log: &LogEntry) -> Vec<ThreatAlert> {
184        let mut alerts = Vec::new();
185
186        // Check source IP
187        if let Some(ref ip) = log.source_ip {
188            if let Some(alert) = self.check_ioc(IOCType::IPAddress, ip, log) {
189                alerts.push(alert);
190            }
191        }
192
193        // Check for domains in message
194        for domain in self.extract_domains(&log.message) {
195            if let Some(alert) = self.check_ioc(IOCType::Domain, &domain, log) {
196                alerts.push(alert);
197            }
198        }
199
200        // Check for file hashes in metadata
201        if let Some(hash) = log.metadata.get("file_hash") {
202            if let Some(alert) = self.check_ioc(IOCType::FileHash, hash, log) {
203                alerts.push(alert);
204            }
205        }
206
207        // Check for URLs in metadata
208        if let Some(url) = log.metadata.get("url") {
209            if let Some(alert) = self.check_ioc(IOCType::URL, url, log) {
210                alerts.push(alert);
211            }
212        }
213
214        // Check user agent
215        if let Some(user_agent) = log.metadata.get("user_agent") {
216            for ioc_map in self.iocs.get(&IOCType::UserAgent).iter() {
217                for (pattern, ioc) in ioc_map.iter() {
218                    if user_agent.to_lowercase().contains(&pattern.to_lowercase()) {
219                        self.matches_count += 1;
220                        alerts.push(self.create_alert(ioc, log));
221                    }
222                }
223            }
224        }
225
226        alerts
227    }
228
229    /// Check specific IOC
230    fn check_ioc(&mut self, ioc_type: IOCType, value: &str, log: &LogEntry) -> Option<ThreatAlert> {
231        if let Some(type_map) = self.iocs.get(&ioc_type) {
232            if let Some(ioc) = type_map.get(value) {
233                self.matches_count += 1;
234                return Some(self.create_alert(ioc, log));
235            }
236        }
237        None
238    }
239
240    /// Create alert from IOC match
241    fn create_alert(&self, ioc: &IOC, log: &LogEntry) -> ThreatAlert {
242        ThreatAlert {
243            alert_id: format!("IOC-{}", self.matches_count),
244            timestamp: Utc::now(),
245            severity: ioc.severity,
246            category: ThreatCategory::SystemCompromise,
247            description: format!(
248                "Threat Intelligence Match: {} ({:?})",
249                ioc.description, ioc.ioc_type
250            ),
251            source_log: format!("{} - {}", log.timestamp, log.message),
252            indicators: vec![
253                format!("{:?}: {}", ioc.ioc_type, ioc.value),
254                format!("Source: {}", ioc.source),
255                format!("Confidence: {:.0}%", ioc.confidence * 100.0),
256            ],
257            recommended_action: format!(
258                "Block {} {}, investigate affected systems, review network traffic",
259                format!("{:?}", ioc.ioc_type).to_lowercase(),
260                ioc.value
261            ),
262            threat_score: self.calculate_threat_score(ioc),
263            correlated_alerts: vec![],
264        }
265    }
266
267    /// Calculate threat score from IOC
268    fn calculate_threat_score(&self, ioc: &IOC) -> u32 {
269        let base_score = match ioc.severity {
270            ThreatSeverity::Info => 10,
271            ThreatSeverity::Low => 25,
272            ThreatSeverity::Medium => 50,
273            ThreatSeverity::High => 75,
274            ThreatSeverity::Critical => 95,
275        };
276
277        let confidence_adjustment = (ioc.confidence * 10.0) as u32;
278        (base_score + confidence_adjustment).min(100)
279    }
280
281    /// Extract domains from text
282    fn extract_domains(&self, text: &str) -> Vec<String> {
283        let mut domains = Vec::new();
284        let words: Vec<&str> = text.split_whitespace().collect();
285
286        for word in words {
287            if word.contains('.') && !word.starts_with("http") {
288                // Simple domain detection
289                if let Some(domain) = word.split('/').next() {
290                    if domain.contains('.') {
291                        domains.push(domain.to_string());
292                    }
293                }
294            }
295        }
296
297        domains
298    }
299
300    /// Get IOC by type and value
301    pub fn get_ioc(&self, ioc_type: IOCType, value: &str) -> Option<&IOC> {
302        self.iocs.get(&ioc_type)?.get(value)
303    }
304
305    /// Get all IOCs of a specific type
306    pub fn get_iocs_by_type(&self, ioc_type: IOCType) -> Vec<&IOC> {
307        self.iocs
308            .get(&ioc_type)
309            .map(|map| map.values().collect())
310            .unwrap_or_default()
311    }
312
313    /// Get statistics
314    pub fn get_stats(&self) -> HashMap<String, usize> {
315        let mut stats = HashMap::new();
316        stats.insert("total_matches".to_string(), self.matches_count);
317        stats.insert("malicious_ips".to_string(), self.malicious_ips.len());
318        stats.insert(
319            "malicious_domains".to_string(),
320            self.malicious_domains.len(),
321        );
322        stats.insert("malicious_hashes".to_string(), self.malicious_hashes.len());
323        stats.insert("threat_actors".to_string(), self.threat_actors.len());
324
325        let total_iocs: usize = self.iocs.values().map(|m| m.len()).sum();
326        stats.insert("total_iocs".to_string(), total_iocs);
327
328        stats
329    }
330
331    /// Import IOCs from JSON
332    pub fn import_iocs_json(&mut self, json: &str) -> Result<usize, serde_json::Error> {
333        let iocs: Vec<IOC> = serde_json::from_str(json)?;
334        let count = iocs.len();
335        for ioc in iocs {
336            self.add_ioc(ioc);
337        }
338        Ok(count)
339    }
340
341    /// Export IOCs to JSON
342    pub fn export_iocs_json(&self) -> Result<String, serde_json::Error> {
343        let all_iocs: Vec<&IOC> = self.iocs.values().flat_map(|m| m.values()).collect();
344        serde_json::to_string_pretty(&all_iocs)
345    }
346
347    /// Clear old IOCs
348    pub fn clear_old_iocs(&mut self, before: DateTime<Utc>) {
349        for type_map in self.iocs.values_mut() {
350            type_map.retain(|_, ioc| ioc.last_seen >= before);
351        }
352        // Rebuild fast lookup sets
353        self.rebuild_lookup_sets();
354    }
355
356    /// Rebuild fast lookup sets
357    fn rebuild_lookup_sets(&mut self) {
358        self.malicious_ips.clear();
359        self.malicious_domains.clear();
360        self.malicious_hashes.clear();
361
362        if let Some(ip_map) = self.iocs.get(&IOCType::IPAddress) {
363            self.malicious_ips.extend(ip_map.keys().cloned());
364        }
365        if let Some(domain_map) = self.iocs.get(&IOCType::Domain) {
366            self.malicious_domains.extend(domain_map.keys().cloned());
367        }
368        if let Some(hash_map) = self.iocs.get(&IOCType::FileHash) {
369            self.malicious_hashes.extend(hash_map.keys().cloned());
370        }
371    }
372}
373
374impl Default for ThreatIntelligence {
375    fn default() -> Self {
376        Self::new()
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383    use std::collections::HashMap;
384
385    fn create_log_with_ip(ip: &str) -> LogEntry {
386        LogEntry {
387            timestamp: Utc::now(),
388            source_ip: Some(ip.to_string()),
389            user: Some("test_user".to_string()),
390            event_type: "connection".to_string(),
391            message: "Network connection established".to_string(),
392            metadata: HashMap::new(),
393        }
394    }
395
396    #[test]
397    fn test_ioc_creation() {
398        let ioc = IOC::new(
399            IOCType::IPAddress,
400            "1.2.3.4".to_string(),
401            ThreatSeverity::High,
402            "Test IOC".to_string(),
403            "TestFeed".to_string(),
404        );
405
406        assert_eq!(ioc.value, "1.2.3.4");
407        assert_eq!(ioc.severity, ThreatSeverity::High);
408    }
409
410    #[test]
411    fn test_add_ioc() {
412        let mut intel = ThreatIntelligence::new();
413        let initial_count = intel.malicious_ips.len();
414
415        intel.add_ioc(IOC::new(
416            IOCType::IPAddress,
417            "10.0.0.1".to_string(),
418            ThreatSeverity::Medium,
419            "Test IP".to_string(),
420            "Test".to_string(),
421        ));
422
423        assert_eq!(intel.malicious_ips.len(), initial_count + 1);
424        assert!(intel.malicious_ips.contains("10.0.0.1"));
425    }
426
427    #[test]
428    fn test_check_malicious_ip() {
429        let mut intel = ThreatIntelligence::new();
430
431        // Add malicious IP
432        intel.add_ioc(IOC::new(
433            IOCType::IPAddress,
434            "99.99.99.99".to_string(),
435            ThreatSeverity::Critical,
436            "Malicious server".to_string(),
437            "ThreatFeed".to_string(),
438        ));
439
440        let log = create_log_with_ip("99.99.99.99");
441        let alerts = intel.check_log(&log);
442
443        assert!(!alerts.is_empty());
444        assert_eq!(alerts[0].severity, ThreatSeverity::Critical);
445        assert!(alerts[0].description.contains("Threat Intelligence Match"));
446    }
447
448    #[test]
449    fn test_check_clean_ip() {
450        let mut intel = ThreatIntelligence::new();
451        let log = create_log_with_ip("192.168.1.1");
452        let alerts = intel.check_log(&log);
453
454        // Should not match any default IOCs
455        assert!(alerts
456            .iter()
457            .all(|a| a.severity != ThreatSeverity::Critical));
458    }
459
460    #[test]
461    fn test_domain_extraction() {
462        let intel = ThreatIntelligence::new();
463        let text = "User accessed malicious-example.com and another.domain.org";
464        let domains = intel.extract_domains(text);
465
466        assert!(domains.contains(&"malicious-example.com".to_string()));
467        assert!(domains.contains(&"another.domain.org".to_string()));
468    }
469
470    #[test]
471    fn test_get_ioc() {
472        let mut intel = ThreatIntelligence::new();
473
474        intel.add_ioc(IOC::new(
475            IOCType::Domain,
476            "evil.com".to_string(),
477            ThreatSeverity::High,
478            "Malicious domain".to_string(),
479            "Test".to_string(),
480        ));
481
482        let ioc = intel.get_ioc(IOCType::Domain, "evil.com");
483        assert!(ioc.is_some());
484        assert_eq!(ioc.unwrap().value, "evil.com");
485    }
486
487    #[test]
488    fn test_get_iocs_by_type() {
489        let intel = ThreatIntelligence::new();
490        let ip_iocs = intel.get_iocs_by_type(IOCType::IPAddress);
491
492        assert!(!ip_iocs.is_empty()); // Should have default IOCs
493    }
494
495    #[test]
496    fn test_stats() {
497        let intel = ThreatIntelligence::new();
498        let stats = intel.get_stats();
499
500        assert!(stats.contains_key("total_iocs"));
501        assert!(stats.contains_key("malicious_ips"));
502        assert!(stats.get("total_iocs").unwrap() > &0); // Has default IOCs
503    }
504
505    #[test]
506    fn test_user_agent_detection() {
507        let mut intel = ThreatIntelligence::new();
508        let mut log = LogEntry {
509            timestamp: Utc::now(),
510            source_ip: Some("192.168.1.1".to_string()),
511            user: Some("attacker".to_string()),
512            event_type: "web_request".to_string(),
513            message: "HTTP request".to_string(),
514            metadata: HashMap::new(),
515        };
516
517        log.metadata
518            .insert("user_agent".to_string(), "sqlmap/1.0".to_string());
519
520        let alerts = intel.check_log(&log);
521        assert!(!alerts.is_empty());
522        assert!(
523            alerts[0].description.contains("UserAgent")
524                || alerts[0].description.contains("SQL injection")
525        );
526    }
527
528    #[test]
529    fn test_threat_score_calculation() {
530        let intel = ThreatIntelligence::new();
531
532        let high_confidence_ioc = IOC {
533            ioc_type: IOCType::IPAddress,
534            value: "1.2.3.4".to_string(),
535            severity: ThreatSeverity::Critical,
536            description: "Test".to_string(),
537            source: "Test".to_string(),
538            first_seen: Utc::now(),
539            last_seen: Utc::now(),
540            confidence: 1.0,
541        };
542
543        let score = intel.calculate_threat_score(&high_confidence_ioc);
544        assert!(score >= 95);
545    }
546
547    #[test]
548    fn test_clear_old_iocs() {
549        let mut intel = ThreatIntelligence::new();
550
551        let old_ioc = IOC {
552            ioc_type: IOCType::IPAddress,
553            value: "1.1.1.1".to_string(),
554            severity: ThreatSeverity::Low,
555            description: "Old IOC".to_string(),
556            source: "Test".to_string(),
557            first_seen: Utc::now() - chrono::Duration::days(100),
558            last_seen: Utc::now() - chrono::Duration::days(100),
559            confidence: 0.5,
560        };
561
562        intel.add_ioc(old_ioc);
563
564        let cutoff = Utc::now() - chrono::Duration::days(50);
565        intel.clear_old_iocs(cutoff);
566
567        // Old IOC should be removed
568        assert!(intel.get_ioc(IOCType::IPAddress, "1.1.1.1").is_none());
569    }
570}
rust_threat_detector/threat_intelligence.rs

rust_threat_detector/
threat_intelligence.rs
