ploit_malper/engine/
dedup.rs

1use std::collections::HashMap;
2
3use regex::Regex;
4use sha2::{Digest, Sha256};
5
6use crate::analyze::extract_endpoint;
7use crate::models::{DedupResult, ScanRecord};
8
9fn normalize_title(title: &str) -> String {
10    let cleaned = title
11        .trim()
12        .to_lowercase()
13        .split_whitespace()
14        .collect::<Vec<&str>>()
15        .join(" ");
16
17    let re = Regex::new(r"\s*[/\\]\s*$").expect("valid regex");
18    re.replace_all(&cleaned, "").to_string()
19}
20
21fn composite_key(record: &ScanRecord) -> String {
22    let normalized = normalize_title(&record.title);
23    let raw = format!("{}|{}|{}", record.target, record.tool, normalized);
24    let mut hasher = Sha256::new();
25    hasher.update(raw.as_bytes());
26    format!("{:x}", hasher.finalize())
27}
28
29pub fn deduplicate_records(records: Vec<ScanRecord>) -> DedupResult {
30    let total_count = records.len();
31    let mut seen: HashMap<String, usize> = HashMap::new();
32    let mut unique_records = Vec::new();
33
34    for record in records {
35        let key = composite_key(&record);
36        if let Some(&existing_idx) = seen.get(&key) {
37            merge_record(&mut unique_records[existing_idx], &record);
38        } else {
39            let mut record = record;
40            if record.affected_endpoints.is_empty() {
41                record.affected_endpoints = extract_endpoint(&record).into_iter().collect();
42            }
43            seen.insert(key, unique_records.len());
44            unique_records.push(record);
45        }
46    }
47
48    let unique_count = unique_records.len();
49    let removed_count = total_count.saturating_sub(unique_count);
50
51    DedupResult {
52        unique_count,
53        removed_count,
54        records: unique_records,
55    }
56}
57
58/// Merge location and CVE information from a duplicate into the kept record
59/// so the same vulnerability reported at multiple endpoints is preserved.
60fn merge_record(kept: &mut ScanRecord, duplicate: &ScanRecord) {
61    if let Some(endpoint) = extract_endpoint(duplicate) {
62        if !kept.affected_endpoints.contains(&endpoint) {
63            kept.affected_endpoints.push(endpoint);
64        }
65    }
66    if kept.cve.is_none() {
67        kept.cve = duplicate.cve.clone();
68    }
69    if kept.service.is_none() {
70        kept.service = duplicate.service.clone();
71    }
72    if kept.port.is_none() {
73        kept.port = duplicate.port;
74    }
75}
76
77pub fn normalize_title_text(title: &str) -> String {
78    normalize_title(title)
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use crate::models::ScanRecord;
85    use serde_json::json;
86
87    #[test]
88    fn deduplicates_on_normalized_title() {
89        let records = vec![
90            ScanRecord {
91                target: "10.0.0.1".into(),
92                tool: "nmap".into(),
93                title: "Apache Version /".into(),
94                severity: "high".into(),
95                port: Some(80),
96                cve: None,
97                service: Some("apache".into()),
98                raw: json!({}),
99                ..Default::default()
100            },
101            ScanRecord {
102                target: "10.0.0.1".into(),
103                tool: "nmap".into(),
104                title: "apache version".into(),
105                severity: "high".into(),
106                port: Some(80),
107                cve: None,
108                service: Some("apache".into()),
109                raw: json!({}),
110                ..Default::default()
111            },
112        ];
113
114        let result = deduplicate_records(records);
115        assert_eq!(result.unique_count, 1);
116        assert_eq!(result.removed_count, 1);
117    }
118}
ploit_malper/engine/dedup.rs

ploit_malper/engine/
dedup.rs
