1use std::collections::HashMap;
2use std::fs;
3use std::process::Command;
4use std::thread;
5use std::time::Duration;
6
7use serde_json::Value;
8
9use crate::config::nvd_cache_file;
10use crate::error::Result;
11use crate::models::{NVDCVEInfo, ScanRecord};
12
13const NVD_API_BASE: &str = "https://services.nvd.nist.gov/rest/json/cves/2.0";
14const NVD_RATE_LIMIT_DELAY: Duration = Duration::from_millis(600);
15
16#[derive(Debug, Clone)]
17pub struct NVDClient {
18 pub api_key: String,
19 cache: HashMap<String, serde_json::Value>,
20}
21
22#[derive(Debug, Clone, Copy, Default)]
23pub struct EnrichmentStats {
24 pub fetched: usize,
25 pub cache_hits: usize,
26 pub not_found: usize,
27}
28
29impl NVDClient {
30 pub fn new(api_key: &str) -> Self {
31 let mut client = Self {
32 api_key: api_key.to_string(),
33 cache: HashMap::new(),
34 };
35 client.load_cache();
36 client
37 }
38
39 fn load_cache(&mut self) {
40 let path = nvd_cache_file();
41 if let Ok(raw) = fs::read_to_string(path) {
42 if let Ok(cache) = serde_json::from_str::<HashMap<String, Value>>(&raw) {
43 self.cache = cache;
44 }
45 }
46 }
47
48 fn save_cache(&self) -> Result<()> {
49 let path = nvd_cache_file();
50 if let Some(parent) = path.parent() {
51 fs::create_dir_all(parent)?;
52 }
53 fs::write(path, serde_json::to_string_pretty(&self.cache)?)?;
54 Ok(())
55 }
56
57 fn get_cached(&self, cve_id: &str) -> Option<&Value> {
58 self.cache.get(cve_id)
59 }
60
61 fn cache_result(&mut self, cve_id: &str, data: Value) -> Result<()> {
62 self.cache.insert(cve_id.to_string(), data);
63 self.save_cache()
64 }
65
66 fn fetch_cve(&self, cve_id: &str, use_key: bool) -> std::result::Result<Value, String> {
67 let url = format!("{}?cveId={}", NVD_API_BASE, cve_id);
68 let mut command = Command::new("curl");
69 command
70 .arg("-sS")
71 .arg("--fail")
72 .arg("--max-time")
73 .arg("15")
74 .arg("-H")
75 .arg("User-Agent: PloitMalper/0.1.0");
76
77 if use_key && !self.api_key.is_empty() {
78 command.arg("-H").arg(format!("apiKey: {}", self.api_key));
79 }
80
81 let output = command.arg(&url).output().map_err(|e| e.to_string())?;
82
83 if output.status.success() {
84 serde_json::from_slice(&output.stdout).map_err(|e| e.to_string())
85 } else if output.status.code() == Some(22) && use_key && !self.api_key.is_empty() {
86 let fallback = Command::new("curl")
87 .arg("-sS")
88 .arg("--fail")
89 .arg("--max-time")
90 .arg("15")
91 .arg("-H")
92 .arg("User-Agent: PloitMalper/0.1.0")
93 .arg(&url)
94 .output()
95 .map_err(|e| e.to_string())?;
96 if fallback.status.success() {
97 serde_json::from_slice(&fallback.stdout).map_err(|e| e.to_string())
98 } else {
99 Err(String::from_utf8_lossy(&fallback.stderr).trim().to_string())
100 }
101 } else {
102 Err(String::from_utf8_lossy(&output.stderr).trim().to_string())
103 }
104 }
105
106 pub fn lookup_cve(&mut self, cve_id: &str) -> Option<NVDCVEInfo> {
107 if let Some(cached) = self.get_cached(cve_id) {
108 return Some(Self::info_from_cached(cve_id, cached));
109 }
110
111 thread::sleep(NVD_RATE_LIMIT_DELAY);
112
113 let resp = match self.fetch_cve(cve_id, true) {
114 Ok(value) => value,
115 Err(_) => return None,
116 };
117
118 let vulnerabilities = resp.get("vulnerabilities")?.as_array()?;
119 let cve_data = vulnerabilities.first()?.get("cve")?.as_object()?.clone();
120
121 let info = Self::info_from_cve(cve_id, &cve_data);
122 let cache_entry = serde_json::json!({
123 "cve_id": info.cve_id,
124 "description": info.description,
125 "cvss_v3_score": info.cvss_v3_score,
126 "cvss_v3_severity": info.cvss_v3_severity,
127 "cvss_v2_score": info.cvss_v2_score,
128 "published": info.published,
129 "last_modified": info.last_modified,
130 "weaknesses": info.weaknesses,
131 "references": info.references,
132 });
133
134 let _ = self.cache_result(cve_id, cache_entry);
135 Some(info)
136 }
137
138 fn info_from_cached(cve_id: &str, cached: &Value) -> NVDCVEInfo {
139 NVDCVEInfo {
140 cve_id: cached
141 .get("cve_id")
142 .and_then(Value::as_str)
143 .unwrap_or(cve_id)
144 .to_string(),
145 description: cached
146 .get("description")
147 .and_then(Value::as_str)
148 .unwrap_or("")
149 .to_string(),
150 cvss_v3_score: cached
151 .get("cvss_v3_score")
152 .and_then(Value::as_f64)
153 .map(|v| v as f32),
154 cvss_v3_severity: cached
155 .get("cvss_v3_severity")
156 .and_then(Value::as_str)
157 .map(ToOwned::to_owned),
158 cvss_v2_score: cached
159 .get("cvss_v2_score")
160 .and_then(Value::as_f64)
161 .map(|v| v as f32),
162 published: cached
163 .get("published")
164 .and_then(Value::as_str)
165 .unwrap_or("")
166 .to_string(),
167 last_modified: cached
168 .get("last_modified")
169 .and_then(Value::as_str)
170 .unwrap_or("")
171 .to_string(),
172 weaknesses: cached
173 .get("weaknesses")
174 .and_then(Value::as_array)
175 .map(|refs| {
176 refs.iter()
177 .filter_map(Value::as_str)
178 .map(ToOwned::to_owned)
179 .collect()
180 })
181 .unwrap_or_default(),
182 references: cached
183 .get("references")
184 .and_then(Value::as_array)
185 .map(|refs| {
186 refs.iter()
187 .filter_map(Value::as_str)
188 .map(ToOwned::to_owned)
189 .collect()
190 })
191 .unwrap_or_default(),
192 }
193 }
194
195 fn info_from_cve(cve_id: &str, cve_data: &serde_json::Map<String, Value>) -> NVDCVEInfo {
196 let descriptions = cve_data
197 .get("descriptions")
198 .and_then(Value::as_array)
199 .cloned()
200 .unwrap_or_default();
201
202 let mut description = String::new();
203 for entry in &descriptions {
204 if entry.get("lang").and_then(Value::as_str) == Some("en") {
205 description = entry
206 .get("value")
207 .and_then(Value::as_str)
208 .unwrap_or("")
209 .to_string();
210 break;
211 }
212 }
213 if description.is_empty() {
214 description = descriptions
215 .first()
216 .and_then(|entry| entry.get("value"))
217 .and_then(Value::as_str)
218 .unwrap_or("")
219 .to_string();
220 }
221
222 let metrics = cve_data.get("metrics").and_then(Value::as_object);
223 let mut cvss_v3_score = None;
224 let mut cvss_v3_severity = None;
225 let mut cvss_v2_score = None;
226
227 if let Some(metrics) = metrics {
228 let metric_groups = metrics
229 .get("cvssMetricV31")
230 .or_else(|| metrics.get("cvssMetricV30"))
231 .and_then(Value::as_array)
232 .cloned()
233 .unwrap_or_default();
234 if let Some(metric_group) = metric_groups.first() {
235 if let Some(cvss_data) = metric_group.get("cvssData").and_then(Value::as_object) {
236 cvss_v3_score = cvss_data
237 .get("baseScore")
238 .and_then(Value::as_f64)
239 .map(|v| v as f32);
240 cvss_v3_severity = cvss_data
241 .get("baseSeverity")
242 .and_then(Value::as_str)
243 .map(ToOwned::to_owned);
244 }
245 }
246
247 if cvss_v3_score.is_none() {
248 if let Some(cvss_v2_metrics) = metrics.get("cvssMetricV2").and_then(Value::as_array)
249 {
250 if let Some(metric) = cvss_v2_metrics.first() {
251 cvss_v2_score = metric
252 .get("cvssData")
253 .and_then(Value::as_object)
254 .and_then(|cvss| cvss.get("baseScore"))
255 .and_then(Value::as_f64)
256 .map(|v| v as f32);
257 }
258 }
259 }
260 }
261
262 let references = cve_data
263 .get("references")
264 .and_then(Value::as_array)
265 .map(|refs| {
266 refs.iter()
267 .take(5)
268 .filter_map(|ref_value| ref_value.get("url").and_then(Value::as_str))
269 .map(ToOwned::to_owned)
270 .collect::<Vec<_>>()
271 })
272 .unwrap_or_default();
273
274 let weaknesses = cve_data
275 .get("weaknesses")
276 .and_then(Value::as_array)
277 .map(|entries| {
278 let mut out = Vec::new();
279 for entry in entries {
280 if let Some(descs) = entry.get("description").and_then(Value::as_array) {
281 for desc in descs {
282 if let Some(value) = desc.get("value").and_then(Value::as_str) {
283 if !out.contains(&value.to_string()) {
284 out.push(value.to_string());
285 }
286 }
287 }
288 }
289 }
290 out
291 })
292 .unwrap_or_default();
293
294 NVDCVEInfo {
295 cve_id: cve_id.to_string(),
296 description,
297 cvss_v3_score,
298 cvss_v3_severity,
299 cvss_v2_score,
300 published: cve_data
301 .get("published")
302 .and_then(Value::as_str)
303 .unwrap_or("")
304 .to_string(),
305 last_modified: cve_data
306 .get("lastModified")
307 .and_then(Value::as_str)
308 .unwrap_or("")
309 .to_string(),
310 weaknesses,
311 references,
312 }
313 }
314
315 pub fn enrich_records(&mut self, records: &mut [ScanRecord]) -> EnrichmentStats {
316 let mut stats = EnrichmentStats::default();
317
318 for record in records {
319 let cves = if record.cves.is_empty() {
320 record.cve.clone().into_iter().collect::<Vec<_>>()
321 } else {
322 record.cves.clone()
323 };
324
325 for cve in cves {
326 let key = cve.trim().to_uppercase();
327 if let Some(cached) = self.get_cached(&key).cloned() {
328 stats.cache_hits += 1;
329 let info = Self::info_from_cached(&key, &cached);
330 apply_nvd(record, &info);
331 continue;
332 }
333
334 if let Some(info) = self.lookup_cve(&key) {
335 stats.fetched += 1;
336 apply_nvd(record, &info);
337 } else {
338 stats.not_found += 1;
339 }
340 }
341 }
342
343 stats
344 }
345}
346
347fn apply_nvd(record: &mut ScanRecord, info: &NVDCVEInfo) {
348 record.nvd_description = Some(info.description.clone());
349 record.nvd_cvss_v3 = info.cvss_v3_score;
350 record.nvd_cvss_v3_severity = info.cvss_v3_severity.clone();
351 record.nvd_cvss_v2 = info.cvss_v2_score;
352 record.nvd_published = Some(info.published.clone());
353 record.nvd_last_modified = Some(info.last_modified.clone());
354 record.nvd_weaknesses = info.weaknesses.clone();
355 record.nvd_references = info.references.clone();
356
357 if let Some(severity) = &info.cvss_v3_severity {
358 record.severity = severity.to_lowercase();
359 }
360}
361
362#[cfg(test)]
363mod tests {
364 use super::*;
365 use serde_json::json;
366
367 #[test]
368 fn parses_cwe_weaknesses_from_nvd_response() {
369 let cve_data = json!({
370 "descriptions": [{"lang": "en", "value": "Path traversal in Apache."}],
371 "metrics": {
372 "cvssMetricV31": [{
373 "cvssData": {
374 "baseScore": 9.8,
375 "baseSeverity": "CRITICAL"
376 }
377 }]
378 },
379 "weaknesses": [{
380 "description": [{"lang": "en", "value": "CWE-22"}, {"lang": "en", "value": "CWE-200"}]
381 }],
382 "published": "2021-10-05T00:00:00.000",
383 "lastModified": "2021-10-08T00:00:00.000",
384 "references": [{"url": "https://example.com/advisory"}]
385 });
386 let info = NVDClient::info_from_cve("CVE-2021-41773", cve_data.as_object().unwrap());
387 assert_eq!(
388 info.weaknesses,
389 vec!["CWE-22".to_string(), "CWE-200".to_string()]
390 );
391 assert_eq!(info.cvss_v3_score, Some(9.8));
392 assert_eq!(info.cvss_v3_severity.as_deref(), Some("CRITICAL"));
393 assert_eq!(info.published, "2021-10-05T00:00:00.000");
394 assert_eq!(
395 info.references,
396 vec!["https://example.com/advisory".to_string()]
397 );
398 }
399
400 #[test]
401 fn enriches_every_discovered_cve_from_cache() {
402 let mut client = NVDClient {
403 api_key: String::new(),
404 cache: HashMap::new(),
405 };
406 client.cache.insert(
407 "CVE-2003-1418".into(),
408 json!({
409 "cve_id": "CVE-2003-1418",
410 "description": "Inode leak",
411 "cvss_v3_score": null,
412 "cvss_v3_severity": null,
413 "cvss_v2_score": 5.0,
414 "published": "2003-08-11T00:00:00.000",
415 "last_modified": "2017-10-11T00:00:00.000",
416 "weaknesses": ["CWE-200"],
417 "references": []
418 }),
419 );
420 client.cache.insert(
421 "CVE-2021-41773".into(),
422 json!({
423 "cve_id": "CVE-2021-41773",
424 "description": "Path traversal",
425 "cvss_v3_score": 9.8,
426 "cvss_v3_severity": "CRITICAL",
427 "cvss_v2_score": null,
428 "published": "2021-10-05T00:00:00.000",
429 "last_modified": "2021-10-08T00:00:00.000",
430 "weaknesses": ["CWE-22"],
431 "references": ["https://example.com"]
432 }),
433 );
434
435 let mut records = vec![ScanRecord {
436 cves: vec!["CVE-2003-1418".to_string(), "CVE-2021-41773".to_string()],
437 ..Default::default()
438 }];
439 let stats = client.enrich_records(&mut records);
440 assert_eq!(stats.cache_hits, 2);
441 assert_eq!(stats.fetched, 0);
442 assert_eq!(stats.not_found, 0);
443 assert_eq!(records[0].nvd_cvss_v3, Some(9.8));
444 assert_eq!(records[0].severity, "critical");
445 assert_eq!(records[0].nvd_weaknesses, vec!["CWE-22".to_string()]);
446 }
447
448 #[test]
449 fn gracefully_handles_missing_cve() {
450 let mut client = NVDClient {
453 api_key: String::new(),
454 cache: HashMap::new(),
455 };
456 client.cache.insert(
457 "CVE-2021-44228".into(),
458 json!({
459 "cve_id": "CVE-2021-44228",
460 "description": "Log4j JNDI injection",
461 "cvss_v3_score": 10.0,
462 "cvss_v3_severity": "CRITICAL",
463 "cvss_v2_score": null,
464 "published": "2021-12-09T00:00:00.000",
465 "last_modified": "2021-12-10T00:00:00.000",
466 "weaknesses": ["CWE-502"],
467 "references": []
468 }),
469 );
470 let mut records = vec![ScanRecord {
471 cves: vec!["CVE-2021-44228".to_string()],
472 ..Default::default()
473 }];
474 let stats = client.enrich_records(&mut records);
475 assert_eq!(stats.cache_hits, 1);
476 assert_eq!(records[0].nvd_cvss_v3, Some(10.0));
477 assert_eq!(records[0].nvd_cvss_v3_severity.as_deref(), Some("CRITICAL"));
478 }
479}