forensic_hashdb/
known_bad.rs

1use crate::types::BadFileInfo;
2#[cfg(test)]
3use crate::types::BadFileSource;
4use std::collections::HashMap;
5
6/// Provenance-tracked malware hash lookup database.
7///
8/// Phase 1: backed by a `HashMap` (exact matching, zero false positives).
9/// Phase 2 (future): `might_be_malicious` will be backed by an XOR filter
10/// pre-screen, but the public API will remain unchanged.
11pub struct KnownBadDb {
12    entries: HashMap<[u8; 32], BadFileInfo>,
13}
14
15impl KnownBadDb {
16    /// Construct from an iterator of `(sha256, info)` pairs.
17    pub fn from_entries(iter: impl IntoIterator<Item = ([u8; 32], BadFileInfo)>) -> Self {
18        Self {
19            entries: iter.into_iter().collect(),
20        }
21    }
22
23    /// Fast pre-screen. Returns `true` if the hash MIGHT be in the database.
24    ///
25    /// No false negatives. Phase 1: backed by `HashMap` (no false positives
26    /// either). Phase 2 will swap in an XOR filter pre-screen.
27    pub fn might_be_malicious(&self, sha256: &[u8; 32]) -> bool {
28        self.entries.contains_key(sha256)
29    }
30
31    /// Exact lookup with full provenance. Returns `None` if definitely not present.
32    pub fn lookup(&self, sha256: &[u8; 32]) -> Option<&BadFileInfo> {
33        self.entries.get(sha256)
34    }
35
36    /// Number of entries in the database.
37    pub fn len(&self) -> usize {
38        self.entries.len()
39    }
40
41    pub fn is_empty(&self) -> bool {
42        self.len() == 0
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    fn make_info(sha256: [u8; 32], source: BadFileSource) -> BadFileInfo {
51        BadFileInfo {
52            sha256,
53            source,
54            malware_family: Some("TestFamily".to_string()),
55            tags: vec!["test".to_string()],
56        }
57    }
58
59    #[test]
60    fn known_bad_lookup_hit_returns_info() {
61        let sha = [0x42u8; 32];
62        let info = make_info(sha, BadFileSource::MalwareBazaar);
63        let db = KnownBadDb::from_entries([(sha, info)]);
64        let result = db.lookup(&sha);
65        assert!(result.is_some());
66        assert_eq!(result.unwrap().sha256, sha);
67    }
68
69    #[test]
70    fn known_bad_lookup_miss_returns_none() {
71        let sha = [0x42u8; 32];
72        let info = make_info(sha, BadFileSource::VirusShare);
73        let db = KnownBadDb::from_entries([(sha, info)]);
74        assert!(db.lookup(&[0x99u8; 32]).is_none());
75    }
76
77    #[test]
78    fn known_bad_might_be_malicious_hit_returns_true() {
79        let sha = [0x11u8; 32];
80        let info = make_info(sha, BadFileSource::Malshare);
81        let db = KnownBadDb::from_entries([(sha, info)]);
82        assert!(db.might_be_malicious(&sha));
83    }
84
85    #[test]
86    fn known_bad_might_be_malicious_miss_returns_false() {
87        let sha = [0x11u8; 32];
88        let info = make_info(sha, BadFileSource::AlienVaultOtx);
89        let db = KnownBadDb::from_entries([(sha, info)]);
90        assert!(!db.might_be_malicious(&[0x22u8; 32]));
91    }
92
93    #[test]
94    fn known_bad_empty_db() {
95        let db = KnownBadDb::from_entries([]);
96        assert!(db.is_empty());
97        assert_eq!(db.len(), 0);
98        assert!(db.lookup(&[0x00u8; 32]).is_none());
99        assert!(!db.might_be_malicious(&[0x00u8; 32]));
100    }
101
102    #[test]
103    fn bad_file_info_source_custom() {
104        let sha = [0xddu8; 32];
105        let info = BadFileInfo {
106            sha256: sha,
107            source: BadFileSource::Custom("my-internal-feed"),
108            malware_family: None,
109            tags: vec![],
110        };
111        let db = KnownBadDb::from_entries([(sha, info)]);
112        let result = db.lookup(&sha).unwrap();
113        assert_eq!(result.source, BadFileSource::Custom("my-internal-feed"));
114    }
115}
forensic_hashdb/known_bad.rs

forensic_hashdb/
known_bad.rs
