1use chrono::{DateTime, Duration, Utc};
7use regex::Regex;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use uuid::Uuid;
11
12use crate::LogEntry;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub enum HuntStatus {
17 Draft,
18 Active,
19 Paused,
20 Completed,
21 Archived,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
26pub enum HuntResultType {
27 NoFindings,
28 FalsePositive,
29 TruePositive,
30 RequiresInvestigation,
31 Inconclusive,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ThreatHunt {
37 pub id: String,
38 pub name: String,
39 pub hypothesis: String,
40 pub description: String,
41 pub status: HuntStatus,
42 pub created_at: DateTime<Utc>,
43 pub started_at: Option<DateTime<Utc>>,
44 pub completed_at: Option<DateTime<Utc>>,
45 pub owner: String,
46 pub mitre_techniques: Vec<String>,
47 pub data_sources: Vec<String>,
48 pub queries: Vec<HuntQuery>,
49 pub findings: Vec<HuntFinding>,
50 pub timeline: Vec<HuntTimelineEntry>,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct HuntQuery {
56 pub id: String,
57 pub name: String,
58 pub description: String,
59 pub query_type: QueryType,
60 pub pattern: String,
61 pub data_source: String,
62 pub expected_results: String,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
67pub enum QueryType {
68 Regex,
69 Keyword,
70 Statistical,
71 Behavioral,
72 IOC,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct HuntFinding {
78 pub id: String,
79 pub timestamp: DateTime<Utc>,
80 pub query_id: String,
81 pub result_type: HuntResultType,
82 pub description: String,
83 pub evidence: Vec<String>,
84 pub affected_assets: Vec<String>,
85 pub severity: FindingSeverity,
86 pub recommendations: Vec<String>,
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
91pub enum FindingSeverity {
92 Informational,
93 Low,
94 Medium,
95 High,
96 Critical,
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct HuntTimelineEntry {
102 pub timestamp: DateTime<Utc>,
103 pub action: String,
104 pub actor: String,
105 pub details: String,
106}
107
108impl ThreatHunt {
109 pub fn new(name: &str, hypothesis: &str, owner: &str) -> Self {
111 Self {
112 id: Uuid::new_v4().to_string(),
113 name: name.to_string(),
114 hypothesis: hypothesis.to_string(),
115 description: String::new(),
116 status: HuntStatus::Draft,
117 created_at: Utc::now(),
118 started_at: None,
119 completed_at: None,
120 owner: owner.to_string(),
121 mitre_techniques: Vec::new(),
122 data_sources: Vec::new(),
123 queries: Vec::new(),
124 findings: Vec::new(),
125 timeline: vec![HuntTimelineEntry {
126 timestamp: Utc::now(),
127 action: "Created".to_string(),
128 actor: owner.to_string(),
129 details: "Hunt created".to_string(),
130 }],
131 }
132 }
133
134 pub fn start(&mut self, actor: &str) {
136 self.status = HuntStatus::Active;
137 self.started_at = Some(Utc::now());
138 self.add_timeline_entry("Started", actor, "Hunt execution started");
139 }
140
141 pub fn pause(&mut self, actor: &str, reason: &str) {
143 self.status = HuntStatus::Paused;
144 self.add_timeline_entry("Paused", actor, reason);
145 }
146
147 pub fn complete(&mut self, actor: &str, summary: &str) {
149 self.status = HuntStatus::Completed;
150 self.completed_at = Some(Utc::now());
151 self.add_timeline_entry("Completed", actor, summary);
152 }
153
154 pub fn add_query(&mut self, query: HuntQuery) {
156 self.queries.push(query);
157 }
158
159 pub fn add_finding(&mut self, finding: HuntFinding) {
161 self.findings.push(finding);
162 }
163
164 pub fn add_timeline_entry(&mut self, action: &str, actor: &str, details: &str) {
166 self.timeline.push(HuntTimelineEntry {
167 timestamp: Utc::now(),
168 action: action.to_string(),
169 actor: actor.to_string(),
170 details: details.to_string(),
171 });
172 }
173
174 pub fn duration(&self) -> Option<Duration> {
176 let start = self.started_at?;
177 let end = self.completed_at.unwrap_or_else(Utc::now);
178 Some(end - start)
179 }
180
181 pub fn count_findings_by_type(&self) -> HashMap<HuntResultType, usize> {
183 let mut counts = HashMap::new();
184 for finding in &self.findings {
185 *counts.entry(finding.result_type).or_insert(0) += 1;
186 }
187 counts
188 }
189}
190
191#[derive(Debug, Clone, Serialize, Deserialize)]
193pub struct HuntIOC {
194 pub indicator: String,
195 pub ioc_type: IOCType,
196 pub description: String,
197 pub confidence: f32,
198 pub source: String,
199 pub tags: Vec<String>,
200}
201
202#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
204pub enum IOCType {
205 IPAddress,
206 Domain,
207 URL,
208 FileHash,
209 FileName,
210 Registry,
211 EmailAddress,
212 UserAgent,
213 ProcessName,
214 Command,
215}
216
217pub struct ThreatHuntingEngine {
219 hunts: HashMap<String, ThreatHunt>,
220 ioc_database: Vec<HuntIOC>,
221 hunt_templates: Vec<HuntTemplate>,
222}
223
224#[derive(Debug, Clone, Serialize, Deserialize)]
226pub struct HuntTemplate {
227 pub id: String,
228 pub name: String,
229 pub description: String,
230 pub hypothesis_template: String,
231 pub mitre_techniques: Vec<String>,
232 pub suggested_queries: Vec<HuntQuery>,
233 pub data_sources: Vec<String>,
234}
235
236impl ThreatHuntingEngine {
237 pub fn new() -> Self {
239 let mut engine = Self {
240 hunts: HashMap::new(),
241 ioc_database: Vec::new(),
242 hunt_templates: Vec::new(),
243 };
244 engine.load_default_templates();
245 engine
246 }
247
248 fn load_default_templates(&mut self) {
250 self.hunt_templates.push(HuntTemplate {
252 id: "TMPL-001".to_string(),
253 name: "Lateral Movement Detection".to_string(),
254 description: "Hunt for signs of lateral movement within the network".to_string(),
255 hypothesis_template: "An attacker may be moving laterally using {{technique}}".to_string(),
256 mitre_techniques: vec!["T1021".to_string(), "T1076".to_string(), "T1077".to_string()],
257 suggested_queries: vec![
258 HuntQuery {
259 id: "Q-001".to_string(),
260 name: "RDP Connections".to_string(),
261 description: "Find unusual RDP connections".to_string(),
262 query_type: QueryType::Regex,
263 pattern: r"(?i)(rdp|3389|remote\s+desktop)".to_string(),
264 data_source: "network_logs".to_string(),
265 expected_results: "RDP connection events".to_string(),
266 },
267 HuntQuery {
268 id: "Q-002".to_string(),
269 name: "PsExec Usage".to_string(),
270 description: "Detect PsExec and similar tools".to_string(),
271 query_type: QueryType::Regex,
272 pattern: r"(?i)(psexec|psexesvc|paexec)".to_string(),
273 data_source: "process_logs".to_string(),
274 expected_results: "PsExec execution events".to_string(),
275 },
276 ],
277 data_sources: vec!["network_logs".to_string(), "auth_logs".to_string(), "process_logs".to_string()],
278 });
279
280 self.hunt_templates.push(HuntTemplate {
282 id: "TMPL-002".to_string(),
283 name: "Credential Theft Detection".to_string(),
284 description: "Hunt for credential dumping and theft activities".to_string(),
285 hypothesis_template: "Credentials may have been compromised via {{method}}".to_string(),
286 mitre_techniques: vec!["T1003".to_string(), "T1110".to_string()],
287 suggested_queries: vec![
288 HuntQuery {
289 id: "Q-003".to_string(),
290 name: "LSASS Access".to_string(),
291 description: "Detect processes accessing LSASS".to_string(),
292 query_type: QueryType::Regex,
293 pattern: r"(?i)(lsass|mimikatz|sekurlsa)".to_string(),
294 data_source: "process_logs".to_string(),
295 expected_results: "LSASS access events".to_string(),
296 },
297 HuntQuery {
298 id: "Q-004".to_string(),
299 name: "Failed Auth Spike".to_string(),
300 description: "Find spikes in failed authentication".to_string(),
301 query_type: QueryType::Statistical,
302 pattern: "failed_auth_count > baseline * 3".to_string(),
303 data_source: "auth_logs".to_string(),
304 expected_results: "Authentication anomalies".to_string(),
305 },
306 ],
307 data_sources: vec!["process_logs".to_string(), "auth_logs".to_string(), "windows_security".to_string()],
308 });
309
310 self.hunt_templates.push(HuntTemplate {
312 id: "TMPL-003".to_string(),
313 name: "Data Exfiltration Detection".to_string(),
314 description: "Hunt for potential data theft and exfiltration".to_string(),
315 hypothesis_template: "Data may be exfiltrating via {{channel}}".to_string(),
316 mitre_techniques: vec!["T1041".to_string(), "T1048".to_string()],
317 suggested_queries: vec![
318 HuntQuery {
319 id: "Q-005".to_string(),
320 name: "Large Outbound Transfers".to_string(),
321 description: "Find unusually large data transfers".to_string(),
322 query_type: QueryType::Statistical,
323 pattern: "bytes_out > 100MB".to_string(),
324 data_source: "network_logs".to_string(),
325 expected_results: "Large outbound transfers".to_string(),
326 },
327 HuntQuery {
328 id: "Q-006".to_string(),
329 name: "DNS Tunneling".to_string(),
330 description: "Detect potential DNS tunneling".to_string(),
331 query_type: QueryType::Behavioral,
332 pattern: "dns_query_length > 50 OR dns_query_entropy > 3.5".to_string(),
333 data_source: "dns_logs".to_string(),
334 expected_results: "Suspicious DNS activity".to_string(),
335 },
336 ],
337 data_sources: vec!["network_logs".to_string(), "dns_logs".to_string(), "proxy_logs".to_string()],
338 });
339
340 self.hunt_templates.push(HuntTemplate {
342 id: "TMPL-004".to_string(),
343 name: "Persistence Mechanism Detection".to_string(),
344 description: "Hunt for attacker persistence mechanisms".to_string(),
345 hypothesis_template: "Attacker may have established persistence using {{mechanism}}".to_string(),
346 mitre_techniques: vec!["T1053".to_string(), "T1547".to_string(), "T1546".to_string()],
347 suggested_queries: vec![
348 HuntQuery {
349 id: "Q-007".to_string(),
350 name: "Scheduled Tasks".to_string(),
351 description: "Find suspicious scheduled tasks".to_string(),
352 query_type: QueryType::Regex,
353 pattern: r"(?i)(schtasks|at\.exe|task\s+scheduler)".to_string(),
354 data_source: "process_logs".to_string(),
355 expected_results: "Scheduled task events".to_string(),
356 },
357 HuntQuery {
358 id: "Q-008".to_string(),
359 name: "Registry Run Keys".to_string(),
360 description: "Detect modifications to run keys".to_string(),
361 query_type: QueryType::Regex,
362 pattern: r"(?i)(run|runonce|userinit)".to_string(),
363 data_source: "registry_logs".to_string(),
364 expected_results: "Registry modifications".to_string(),
365 },
366 ],
367 data_sources: vec!["process_logs".to_string(), "registry_logs".to_string(), "file_logs".to_string()],
368 });
369 }
370
371 pub fn create_hunt_from_template(&mut self, template_id: &str, owner: &str, customization: HashMap<String, String>) -> Option<String> {
373 let template = self.hunt_templates.iter().find(|t| t.id == template_id)?.clone();
374
375 let mut hypothesis = template.hypothesis_template.clone();
376 for (key, value) in &customization {
377 hypothesis = hypothesis.replace(&format!("{{{{{}}}}}", key), value);
378 }
379
380 let mut hunt = ThreatHunt::new(&template.name, &hypothesis, owner);
381 hunt.description = template.description.clone();
382 hunt.mitre_techniques = template.mitre_techniques.clone();
383 hunt.data_sources = template.data_sources.clone();
384
385 for query in &template.suggested_queries {
386 hunt.add_query(query.clone());
387 }
388
389 let id = hunt.id.clone();
390 self.hunts.insert(id.clone(), hunt);
391 Some(id)
392 }
393
394 pub fn create_custom_hunt(&mut self, name: &str, hypothesis: &str, owner: &str) -> String {
396 let hunt = ThreatHunt::new(name, hypothesis, owner);
397 let id = hunt.id.clone();
398 self.hunts.insert(id.clone(), hunt);
399 id
400 }
401
402 pub fn get_hunt(&self, id: &str) -> Option<&ThreatHunt> {
404 self.hunts.get(id)
405 }
406
407 pub fn get_hunt_mut(&mut self, id: &str) -> Option<&mut ThreatHunt> {
409 self.hunts.get_mut(id)
410 }
411
412 pub fn execute_query(&self, query: &HuntQuery, logs: &[LogEntry]) -> Vec<QueryMatch> {
414 let mut matches = Vec::new();
415
416 match query.query_type {
417 QueryType::Regex | QueryType::Keyword => {
418 if let Ok(regex) = Regex::new(&query.pattern) {
419 for log in logs {
420 if regex.is_match(&log.message) {
421 matches.push(QueryMatch {
422 query_id: query.id.clone(),
423 log_timestamp: log.timestamp,
424 matched_content: log.message.clone(),
425 source_ip: log.source_ip.clone(),
426 user: log.user.clone(),
427 match_details: "Regex match".to_string(),
428 });
429 }
430 }
431 }
432 }
433 QueryType::IOC => {
434 for log in logs {
435 for ioc in &self.ioc_database {
436 if log.message.contains(&ioc.indicator) {
437 matches.push(QueryMatch {
438 query_id: query.id.clone(),
439 log_timestamp: log.timestamp,
440 matched_content: log.message.clone(),
441 source_ip: log.source_ip.clone(),
442 user: log.user.clone(),
443 match_details: format!("IOC match: {}", ioc.indicator),
444 });
445 }
446 }
447 }
448 }
449 _ => {
450 }
452 }
453
454 matches
455 }
456
457 pub fn ioc_sweep(&self, logs: &[LogEntry]) -> Vec<IOCSweepResult> {
459 let mut results = Vec::new();
460
461 for ioc in &self.ioc_database {
462 let mut matches = Vec::new();
463
464 for log in logs {
465 if log.message.contains(&ioc.indicator) {
466 matches.push(log.clone());
467 }
468 }
469
470 if !matches.is_empty() {
471 results.push(IOCSweepResult {
472 ioc: ioc.clone(),
473 match_count: matches.len(),
474 first_seen: matches.iter().map(|l| l.timestamp).min(),
475 last_seen: matches.iter().map(|l| l.timestamp).max(),
476 affected_assets: matches
477 .iter()
478 .filter_map(|l| l.source_ip.clone())
479 .collect(),
480 });
481 }
482 }
483
484 results
485 }
486
487 pub fn add_ioc(&mut self, ioc: HuntIOC) {
489 self.ioc_database.push(ioc);
490 }
491
492 pub fn add_iocs(&mut self, iocs: Vec<HuntIOC>) {
494 self.ioc_database.extend(iocs);
495 }
496
497 pub fn get_active_hunts(&self) -> Vec<&ThreatHunt> {
499 self.hunts
500 .values()
501 .filter(|h| h.status == HuntStatus::Active)
502 .collect()
503 }
504
505 pub fn get_templates(&self) -> &[HuntTemplate] {
507 &self.hunt_templates
508 }
509
510 pub fn get_statistics(&self) -> HuntStatistics {
512 let total = self.hunts.len();
513 let active = self.hunts.values().filter(|h| h.status == HuntStatus::Active).count();
514 let completed = self.hunts.values().filter(|h| h.status == HuntStatus::Completed).count();
515
516 let total_findings: usize = self.hunts.values().map(|h| h.findings.len()).sum();
517 let true_positives = self.hunts
518 .values()
519 .flat_map(|h| &h.findings)
520 .filter(|f| f.result_type == HuntResultType::TruePositive)
521 .count();
522
523 HuntStatistics {
524 total_hunts: total,
525 active_hunts: active,
526 completed_hunts: completed,
527 total_findings,
528 true_positives,
529 iocs_in_database: self.ioc_database.len(),
530 templates_available: self.hunt_templates.len(),
531 }
532 }
533}
534
535impl Default for ThreatHuntingEngine {
536 fn default() -> Self {
537 Self::new()
538 }
539}
540
541#[derive(Debug, Clone, Serialize, Deserialize)]
543pub struct QueryMatch {
544 pub query_id: String,
545 pub log_timestamp: DateTime<Utc>,
546 pub matched_content: String,
547 pub source_ip: Option<String>,
548 pub user: Option<String>,
549 pub match_details: String,
550}
551
552#[derive(Debug, Clone, Serialize, Deserialize)]
554pub struct IOCSweepResult {
555 pub ioc: HuntIOC,
556 pub match_count: usize,
557 pub first_seen: Option<DateTime<Utc>>,
558 pub last_seen: Option<DateTime<Utc>>,
559 pub affected_assets: Vec<String>,
560}
561
562#[derive(Debug, Clone, Serialize, Deserialize)]
564pub struct HuntStatistics {
565 pub total_hunts: usize,
566 pub active_hunts: usize,
567 pub completed_hunts: usize,
568 pub total_findings: usize,
569 pub true_positives: usize,
570 pub iocs_in_database: usize,
571 pub templates_available: usize,
572}
573
574#[cfg(test)]
575mod tests {
576 use super::*;
577
578 #[test]
579 fn test_hunt_creation() {
580 let hunt = ThreatHunt::new(
581 "Test Hunt",
582 "An attacker may be present",
583 "analyst",
584 );
585
586 assert_eq!(hunt.status, HuntStatus::Draft);
587 assert_eq!(hunt.owner, "analyst");
588 assert!(hunt.started_at.is_none());
589 }
590
591 #[test]
592 fn test_hunt_lifecycle() {
593 let mut hunt = ThreatHunt::new("Test Hunt", "Hypothesis", "analyst");
594
595 hunt.start("analyst");
596 assert_eq!(hunt.status, HuntStatus::Active);
597 assert!(hunt.started_at.is_some());
598
599 hunt.pause("analyst", "Need more data");
600 assert_eq!(hunt.status, HuntStatus::Paused);
601
602 hunt.start("analyst");
603 hunt.complete("analyst", "No threats found");
604 assert_eq!(hunt.status, HuntStatus::Completed);
605 assert!(hunt.completed_at.is_some());
606 }
607
608 #[test]
609 fn test_hunt_from_template() {
610 let mut engine = ThreatHuntingEngine::new();
611
612 let mut customization = HashMap::new();
613 customization.insert("technique".to_string(), "RDP".to_string());
614
615 let hunt_id = engine.create_hunt_from_template("TMPL-001", "analyst", customization);
616 assert!(hunt_id.is_some());
617
618 let hunt = engine.get_hunt(&hunt_id.unwrap());
619 assert!(hunt.is_some());
620 assert!(hunt.unwrap().hypothesis.contains("RDP"));
621 }
622
623 #[test]
624 fn test_query_execution() {
625 let engine = ThreatHuntingEngine::new();
626
627 let query = HuntQuery {
628 id: "TEST-Q1".to_string(),
629 name: "Test Query".to_string(),
630 description: "Test".to_string(),
631 query_type: QueryType::Regex,
632 pattern: r"(?i)failed\s+login".to_string(),
633 data_source: "auth_logs".to_string(),
634 expected_results: "Failed logins".to_string(),
635 };
636
637 let logs = vec![
638 LogEntry {
639 timestamp: Utc::now(),
640 source_ip: Some("192.168.1.1".to_string()),
641 user: Some("user1".to_string()),
642 event_type: "auth".to_string(),
643 message: "Failed login attempt for admin".to_string(),
644 metadata: HashMap::new(),
645 },
646 LogEntry {
647 timestamp: Utc::now(),
648 source_ip: Some("192.168.1.2".to_string()),
649 user: Some("user2".to_string()),
650 event_type: "auth".to_string(),
651 message: "Successful login".to_string(),
652 metadata: HashMap::new(),
653 },
654 ];
655
656 let matches = engine.execute_query(&query, &logs);
657 assert_eq!(matches.len(), 1);
658 assert!(matches[0].matched_content.contains("Failed login"));
659 }
660
661 #[test]
662 fn test_ioc_sweep() {
663 let mut engine = ThreatHuntingEngine::new();
664
665 engine.add_ioc(HuntIOC {
666 indicator: "evil.com".to_string(),
667 ioc_type: IOCType::Domain,
668 description: "Known malicious domain".to_string(),
669 confidence: 0.9,
670 source: "ThreatFeed".to_string(),
671 tags: vec!["malware".to_string()],
672 });
673
674 let logs = vec![
675 LogEntry {
676 timestamp: Utc::now(),
677 source_ip: Some("192.168.1.1".to_string()),
678 user: None,
679 event_type: "dns".to_string(),
680 message: "DNS query for evil.com".to_string(),
681 metadata: HashMap::new(),
682 },
683 ];
684
685 let results = engine.ioc_sweep(&logs);
686 assert_eq!(results.len(), 1);
687 assert_eq!(results[0].match_count, 1);
688 }
689
690 #[test]
691 fn test_hunt_statistics() {
692 let mut engine = ThreatHuntingEngine::new();
693
694 engine.create_custom_hunt("Hunt 1", "Hypothesis 1", "analyst");
695 engine.create_custom_hunt("Hunt 2", "Hypothesis 2", "analyst");
696
697 let stats = engine.get_statistics();
698 assert_eq!(stats.total_hunts, 2);
699 assert!(stats.templates_available > 0);
700 }
701
702 #[test]
703 fn test_hunt_findings() {
704 let mut hunt = ThreatHunt::new("Test", "Hypothesis", "analyst");
705
706 hunt.add_finding(HuntFinding {
707 id: "F-001".to_string(),
708 timestamp: Utc::now(),
709 query_id: "Q-001".to_string(),
710 result_type: HuntResultType::TruePositive,
711 description: "Found malicious activity".to_string(),
712 evidence: vec!["log1".to_string()],
713 affected_assets: vec!["host1".to_string()],
714 severity: FindingSeverity::High,
715 recommendations: vec!["Investigate".to_string()],
716 });
717
718 hunt.add_finding(HuntFinding {
719 id: "F-002".to_string(),
720 timestamp: Utc::now(),
721 query_id: "Q-001".to_string(),
722 result_type: HuntResultType::FalsePositive,
723 description: "Not malicious".to_string(),
724 evidence: vec![],
725 affected_assets: vec![],
726 severity: FindingSeverity::Informational,
727 recommendations: vec![],
728 });
729
730 let counts = hunt.count_findings_by_type();
731 assert_eq!(counts.get(&HuntResultType::TruePositive), Some(&1));
732 assert_eq!(counts.get(&HuntResultType::FalsePositive), Some(&1));
733 }
734}