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 "references": info.references,
131 });
132
133 let _ = self.cache_result(cve_id, cache_entry);
134 Some(info)
135 }
136
137 fn info_from_cached(cve_id: &str, cached: &Value) -> NVDCVEInfo {
138 NVDCVEInfo {
139 cve_id: cached
140 .get("cve_id")
141 .and_then(Value::as_str)
142 .unwrap_or(cve_id)
143 .to_string(),
144 description: cached
145 .get("description")
146 .and_then(Value::as_str)
147 .unwrap_or("")
148 .to_string(),
149 cvss_v3_score: cached
150 .get("cvss_v3_score")
151 .and_then(Value::as_f64)
152 .map(|v| v as f32),
153 cvss_v3_severity: cached
154 .get("cvss_v3_severity")
155 .and_then(Value::as_str)
156 .map(ToOwned::to_owned),
157 cvss_v2_score: cached
158 .get("cvss_v2_score")
159 .and_then(Value::as_f64)
160 .map(|v| v as f32),
161 published: cached
162 .get("published")
163 .and_then(Value::as_str)
164 .unwrap_or("")
165 .to_string(),
166 last_modified: cached
167 .get("last_modified")
168 .and_then(Value::as_str)
169 .unwrap_or("")
170 .to_string(),
171 references: cached
172 .get("references")
173 .and_then(Value::as_array)
174 .map(|refs| {
175 refs.iter()
176 .filter_map(Value::as_str)
177 .map(ToOwned::to_owned)
178 .collect()
179 })
180 .unwrap_or_default(),
181 }
182 }
183
184 fn info_from_cve(cve_id: &str, cve_data: &serde_json::Map<String, Value>) -> NVDCVEInfo {
185 let descriptions = cve_data
186 .get("descriptions")
187 .and_then(Value::as_array)
188 .cloned()
189 .unwrap_or_default();
190
191 let mut description = String::new();
192 for entry in &descriptions {
193 if entry.get("lang").and_then(Value::as_str) == Some("en") {
194 description = entry
195 .get("value")
196 .and_then(Value::as_str)
197 .unwrap_or("")
198 .to_string();
199 break;
200 }
201 }
202 if description.is_empty() {
203 description = descriptions
204 .first()
205 .and_then(|entry| entry.get("value"))
206 .and_then(Value::as_str)
207 .unwrap_or("")
208 .to_string();
209 }
210
211 let metrics = cve_data.get("metrics").and_then(Value::as_object);
212 let mut cvss_v3_score = None;
213 let mut cvss_v3_severity = None;
214 let mut cvss_v2_score = None;
215
216 if let Some(metrics) = metrics {
217 let metric_groups = metrics
218 .get("cvssMetricV31")
219 .or_else(|| metrics.get("cvssMetricV30"))
220 .and_then(Value::as_array)
221 .cloned()
222 .unwrap_or_default();
223 if let Some(metric_group) = metric_groups.first() {
224 if let Some(cvss_data) = metric_group.get("cvssData").and_then(Value::as_object) {
225 cvss_v3_score = cvss_data
226 .get("baseScore")
227 .and_then(Value::as_f64)
228 .map(|v| v as f32);
229 cvss_v3_severity = cvss_data
230 .get("baseSeverity")
231 .and_then(Value::as_str)
232 .map(ToOwned::to_owned);
233 }
234 }
235
236 if cvss_v3_score.is_none() {
237 if let Some(cvss_v2_metrics) = metrics.get("cvssMetricV2").and_then(Value::as_array)
238 {
239 if let Some(metric) = cvss_v2_metrics.first() {
240 cvss_v2_score = metric
241 .get("cvssData")
242 .and_then(Value::as_object)
243 .and_then(|cvss| cvss.get("baseScore"))
244 .and_then(Value::as_f64)
245 .map(|v| v as f32);
246 }
247 }
248 }
249 }
250
251 let references = cve_data
252 .get("references")
253 .and_then(Value::as_array)
254 .map(|refs| {
255 refs.iter()
256 .take(5)
257 .filter_map(|ref_value| ref_value.get("url").and_then(Value::as_str))
258 .map(ToOwned::to_owned)
259 .collect::<Vec<_>>()
260 })
261 .unwrap_or_default();
262
263 NVDCVEInfo {
264 cve_id: cve_id.to_string(),
265 description,
266 cvss_v3_score,
267 cvss_v3_severity,
268 cvss_v2_score,
269 published: cve_data
270 .get("published")
271 .and_then(Value::as_str)
272 .unwrap_or("")
273 .to_string(),
274 last_modified: cve_data
275 .get("lastModified")
276 .and_then(Value::as_str)
277 .unwrap_or("")
278 .to_string(),
279 references,
280 }
281 }
282
283 pub fn enrich_records(&mut self, records: &mut [ScanRecord]) -> EnrichmentStats {
284 let mut stats = EnrichmentStats::default();
285
286 for record in records {
287 let cve = match record.cve.clone() {
288 Some(cve) => cve,
289 None => continue,
290 };
291
292 if let Some(cached) = self.get_cached(&cve).cloned() {
293 stats.cache_hits += 1;
294 let info = Self::info_from_cached(&cve, &cached);
295 apply_nvd(record, &info);
296 continue;
297 }
298
299 if let Some(info) = self.lookup_cve(&cve) {
300 stats.fetched += 1;
301 apply_nvd(record, &info);
302 } else {
303 stats.not_found += 1;
304 }
305 }
306
307 stats
308 }
309}
310
311fn apply_nvd(record: &mut ScanRecord, info: &NVDCVEInfo) {
312 record.nvd_description = Some(info.description.clone());
313 record.nvd_cvss_v3 = info.cvss_v3_score;
314 record.nvd_cvss_v3_severity = info.cvss_v3_severity.clone();
315 record.nvd_cvss_v2 = info.cvss_v2_score;
316 record.nvd_published = Some(info.published.clone());
317 record.nvd_last_modified = Some(info.last_modified.clone());
318 record.nvd_references = info.references.clone();
319
320 if let Some(severity) = &info.cvss_v3_severity {
321 record.severity = severity.to_lowercase();
322 }
323}