1use std::path::Path;
2
3use rusqlite::{params, Connection, OptionalExtension, Row};
4use serde_json::Value;
5
6use crate::db::models::{Asset, ExploitExecution, Finding, Observation, ScanRun, Service};
7use crate::db::{collections, Storage};
8use crate::error::Result;
9
10pub struct SqliteStorage {
13 conn: Connection,
14}
15
16impl SqliteStorage {
17 pub fn open(path: &str) -> Result<Self> {
18 if path != ":memory:" {
19 if let Some(parent) = Path::new(path).parent() {
20 std::fs::create_dir_all(parent)?;
21 }
22 }
23 let conn = Connection::open(path)?;
24 let storage = Self { conn };
25 storage.create_tables()?;
26 Ok(storage)
27 }
28
29 fn create_tables(&self) -> Result<()> {
30 self.conn.execute_batch(
31 "PRAGMA foreign_keys = ON;
32 CREATE TABLE IF NOT EXISTS Assets (
33 stable_id TEXT PRIMARY KEY,
34 asset_type TEXT,
35 name TEXT NOT NULL,
36 ip TEXT,
37 fqdn TEXT,
38 reverse_dns TEXT,
39 first_seen TEXT,
40 last_seen TEXT,
41 status TEXT,
42 metadata TEXT
43 );
44 CREATE TABLE IF NOT EXISTS Services (
45 stable_id TEXT PRIMARY KEY,
46 asset_id TEXT NOT NULL,
47 port INTEGER,
48 protocol TEXT,
49 service_name TEXT,
50 product TEXT,
51 version TEXT,
52 version_str TEXT,
53 banner TEXT,
54 technologies TEXT,
55 cpes TEXT,
56 first_seen TEXT,
57 last_seen TEXT,
58 status TEXT,
59 metadata TEXT
60 );
61 CREATE TABLE IF NOT EXISTS Findings (
62 stable_id TEXT PRIMARY KEY,
63 asset_id TEXT NOT NULL,
64 service_id TEXT,
65 title TEXT NOT NULL,
66 severity TEXT,
67 tool TEXT,
68 target_url TEXT,
69 detail TEXT,
70 detail_path TEXT,
71 reference TEXT,
72 cves TEXT,
73 endpoints TEXT,
74 technologies TEXT,
75 first_seen TEXT,
76 last_seen TEXT,
77 status TEXT,
78 exploitability TEXT,
79 metadata TEXT
80 );
81 CREATE TABLE IF NOT EXISTS Observations (
82 stable_id TEXT PRIMARY KEY,
83 run_id TEXT NOT NULL,
84 subject_type TEXT NOT NULL,
85 subject_id TEXT NOT NULL,
86 kind TEXT NOT NULL,
87 before TEXT,
88 after TEXT,
89 detail TEXT,
90 observed_at TEXT
91 );
92 CREATE TABLE IF NOT EXISTS ScanRuns (
93 stable_id TEXT PRIMARY KEY,
94 target TEXT NOT NULL,
95 folder TEXT,
96 started_at TEXT,
97 finished_at TEXT,
98 tools TEXT,
99 artifacts TEXT,
100 content_hash TEXT,
101 imported_at TEXT,
102 stats TEXT
103 );
104 CREATE TABLE IF NOT EXISTS ExploitExecutions (
105 execution_id TEXT PRIMARY KEY,
106 run_id TEXT NOT NULL,
107 asset_id TEXT NOT NULL,
108 vulnerability_id TEXT NOT NULL,
109 module_type TEXT,
110 module TEXT NOT NULL,
111 host TEXT,
112 payload TEXT,
113 status TEXT,
114 start_time TEXT,
115 finish_time TEXT,
116 job_id TEXT,
117 session_id TEXT,
118 error TEXT,
119 loot TEXT,
120 options TEXT,
121 selected INTEGER,
122 created_at TEXT
123 );
124 CREATE INDEX IF NOT EXISTS idx_services_asset ON Services(asset_id);
125 CREATE INDEX IF NOT EXISTS idx_findings_asset ON Findings(asset_id);
126 CREATE INDEX IF NOT EXISTS idx_obs_subject ON Observations(subject_type, subject_id);
127 CREATE INDEX IF NOT EXISTS idx_exploit_run ON ExploitExecutions(run_id);",
128 )?;
129 Ok(())
130 }
131}
132
133fn json_or_null<T: serde::Serialize>(value: &T) -> rusqlite::Result<String> {
134 serde_json::to_string(value).map_err(|e| rusqlite::Error::ToSqlConversionFailure(e.into()))
135}
136
137fn parse_json(raw: Option<String>) -> Value {
138 raw.and_then(|s| serde_json::from_str(&s).ok())
139 .unwrap_or(Value::Null)
140}
141
142fn parse_vec(raw: Option<String>) -> Vec<String> {
143 match parse_json(raw) {
144 Value::Array(items) => items
145 .into_iter()
146 .filter_map(|v| v.as_str().map(str::to_string))
147 .collect(),
148 _ => Vec::new(),
149 }
150}
151
152fn parse_vec_value(raw: Option<String>) -> Vec<Value> {
153 match parse_json(raw) {
154 Value::Array(items) => items,
155 _ => Vec::new(),
156 }
157}
158
159fn asset_from_row(row: &Row) -> rusqlite::Result<Asset> {
164 Ok(Asset {
165 stable_id: row.get("stable_id")?,
166 asset_type: row.get("asset_type")?,
167 name: row.get("name")?,
168 ip: row.get("ip")?,
169 fqdn: row.get("fqdn")?,
170 reverse_dns: row.get("reverse_dns")?,
171 first_seen: row.get("first_seen")?,
172 last_seen: row.get("last_seen")?,
173 status: row.get("status")?,
174 metadata: parse_json(row.get("metadata")?),
175 })
176}
177
178fn service_from_row(row: &Row) -> rusqlite::Result<Service> {
179 Ok(Service {
180 stable_id: row.get("stable_id")?,
181 asset_id: row.get("asset_id")?,
182 port: row.get("port")?,
183 protocol: row.get("protocol")?,
184 service_name: row.get("service_name")?,
185 product: row.get("product")?,
186 version: row.get("version")?,
187 version_str: row.get("version_str")?,
188 banner: row.get("banner")?,
189 technologies: parse_vec(row.get("technologies")?),
190 cpes: parse_vec(row.get("cpes")?),
191 first_seen: row.get("first_seen")?,
192 last_seen: row.get("last_seen")?,
193 status: row.get("status")?,
194 metadata: parse_json(row.get("metadata")?),
195 })
196}
197
198fn finding_from_row(row: &Row) -> rusqlite::Result<Finding> {
199 Ok(Finding {
200 stable_id: row.get("stable_id")?,
201 asset_id: row.get("asset_id")?,
202 service_id: row.get("service_id")?,
203 title: row.get("title")?,
204 severity: row.get("severity")?,
205 tool: row.get("tool")?,
206 target_url: row.get("target_url")?,
207 detail: row.get("detail")?,
208 detail_path: row.get("detail_path")?,
209 reference: row.get("reference")?,
210 cves: parse_vec(row.get("cves")?),
211 endpoints: parse_vec(row.get("endpoints")?),
212 technologies: parse_vec(row.get("technologies")?),
213 first_seen: row.get("first_seen")?,
214 last_seen: row.get("last_seen")?,
215 status: row.get("status")?,
216 exploitability: row.get("exploitability")?,
217 metadata: parse_json(row.get("metadata")?),
218 })
219}
220
221fn observation_from_row(row: &Row) -> rusqlite::Result<Observation> {
222 Ok(Observation {
223 stable_id: row.get("stable_id")?,
224 run_id: row.get("run_id")?,
225 subject_type: row.get("subject_type")?,
226 subject_id: row.get("subject_id")?,
227 kind: row.get("kind")?,
228 before: parse_json(row.get("before")?),
229 after: parse_json(row.get("after")?),
230 detail: row.get("detail")?,
231 observed_at: row.get("observed_at")?,
232 })
233}
234
235fn scan_run_from_row(row: &Row) -> rusqlite::Result<ScanRun> {
236 Ok(ScanRun {
237 stable_id: row.get("stable_id")?,
238 target: row.get("target")?,
239 folder: row.get("folder")?,
240 started_at: row.get("started_at")?,
241 finished_at: row.get("finished_at")?,
242 tools: parse_vec(row.get("tools")?),
243 artifacts: parse_vec_value(row.get("artifacts")?),
244 content_hash: row.get("content_hash")?,
245 imported_at: row.get("imported_at")?,
246 stats: parse_json(row.get("stats")?),
247 })
248}
249
250fn exploit_execution_from_row(row: &Row) -> rusqlite::Result<ExploitExecution> {
251 Ok(ExploitExecution {
252 execution_id: row.get("execution_id")?,
253 run_id: row.get("run_id")?,
254 asset_id: row.get("asset_id")?,
255 vulnerability_id: row.get("vulnerability_id")?,
256 module_type: row.get("module_type")?,
257 module: row.get("module")?,
258 host: row.get("host")?,
259 payload: row.get("payload")?,
260 status: row.get("status")?,
261 start_time: row.get("start_time")?,
262 finish_time: row.get("finish_time")?,
263 job_id: row.get("job_id")?,
264 session_id: row.get("session_id")?,
265 error: row.get("error")?,
266 loot: parse_vec(row.get("loot")?),
267 options: parse_options_map(row.get("options")?),
268 selected: row.get("selected")?,
269 created_at: row.get("created_at")?,
270 })
271}
272
273fn parse_options_map(raw: Option<String>) -> std::collections::BTreeMap<String, String> {
274 match parse_json(raw) {
275 Value::Object(map) => map
276 .into_iter()
277 .filter_map(|(key, val)| val.as_str().map(|s| (key, s.to_string())))
278 .collect(),
279 _ => std::collections::BTreeMap::new(),
280 }
281}
282
283impl Storage for SqliteStorage {
284 fn kind(&self) -> &str {
285 "sqlite"
286 }
287
288 fn ping(&mut self) -> Result<bool> {
289 let value: String = self
290 .conn
291 .query_row(
292 "SELECT name FROM sqlite_master WHERE type='table' AND name=?1 LIMIT 1",
293 params![collections::ASSETS],
294 |row| row.get(0),
295 )
296 .optional()?
297 .unwrap_or_default();
298 Ok(value == collections::ASSETS)
299 }
300
301 fn ensure_schema(&mut self) -> Result<()> {
302 self.create_tables()
303 }
304
305 fn upsert_scan_run(&mut self, run: &ScanRun) -> Result<()> {
306 self.conn.execute(
307 "INSERT OR REPLACE INTO ScanRuns (stable_id, target, folder, started_at, finished_at, tools, artifacts, content_hash, imported_at, stats)
308 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
309 params![
310 run.stable_id,
311 run.target,
312 run.folder,
313 run.started_at,
314 run.finished_at,
315 json_or_null(&run.tools)?,
316 json_or_null(&run.artifacts)?,
317 run.content_hash,
318 run.imported_at,
319 json_or_null(&run.stats)?
320 ],
321 )?;
322 Ok(())
323 }
324
325 fn get_scan_run(&mut self, run_id: &str) -> Result<Option<ScanRun>> {
326 let row = self
327 .conn
328 .query_row(
329 "SELECT * FROM ScanRuns WHERE stable_id = ?1",
330 params![run_id],
331 scan_run_from_row,
332 )
333 .optional()?;
334 Ok(row)
335 }
336
337 fn list_scan_runs(&mut self) -> Result<Vec<ScanRun>> {
338 let mut stmt = self
339 .conn
340 .prepare("SELECT * FROM ScanRuns ORDER BY imported_at DESC")?;
341 let rows = stmt.query_map([], scan_run_from_row)?;
342 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
343 }
344
345 fn delete_scan_run(&mut self, run_id: &str) -> Result<()> {
346 self.conn.execute(
347 "DELETE FROM ExploitExecutions WHERE run_id = ?1",
348 params![run_id],
349 )?;
350 self.conn.execute(
351 "DELETE FROM Observations WHERE run_id = ?1",
352 params![run_id],
353 )?;
354 self.conn
355 .execute("DELETE FROM ScanRuns WHERE stable_id = ?1", params![run_id])?;
356 Ok(())
357 }
358
359 fn upsert_asset(&mut self, asset: &Asset) -> Result<()> {
360 self.conn.execute(
361 "INSERT OR REPLACE INTO Assets (stable_id, asset_type, name, ip, fqdn, reverse_dns, first_seen, last_seen, status, metadata)
362 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
363 params![
364 asset.stable_id,
365 asset.asset_type,
366 asset.name,
367 asset.ip,
368 asset.fqdn,
369 asset.reverse_dns,
370 asset.first_seen,
371 asset.last_seen,
372 asset.status,
373 json_or_null(&asset.metadata)?
374 ],
375 )?;
376 Ok(())
377 }
378
379 fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>> {
380 let row = self
381 .conn
382 .query_row(
383 "SELECT * FROM Assets WHERE stable_id = ?1",
384 params![asset_id],
385 asset_from_row,
386 )
387 .optional()?;
388 Ok(row)
389 }
390
391 fn list_assets(&mut self) -> Result<Vec<Asset>> {
392 let mut stmt = self.conn.prepare("SELECT * FROM Assets ORDER BY name")?;
393 let rows = stmt.query_map([], asset_from_row)?;
394 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
395 }
396
397 fn upsert_service(&mut self, service: &Service) -> Result<()> {
398 self.conn.execute(
399 "INSERT OR REPLACE INTO Services (stable_id, asset_id, port, protocol, service_name, product, version, version_str, banner, technologies, cpes, first_seen, last_seen, status, metadata)
400 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)",
401 params![
402 service.stable_id,
403 service.asset_id,
404 service.port,
405 service.protocol,
406 service.service_name,
407 service.product,
408 service.version,
409 service.version_str,
410 service.banner,
411 json_or_null(&service.technologies)?,
412 json_or_null(&service.cpes)?,
413 service.first_seen,
414 service.last_seen,
415 service.status,
416 json_or_null(&service.metadata)?
417 ],
418 )?;
419 Ok(())
420 }
421
422 fn get_service(&mut self, service_id: &str) -> Result<Option<Service>> {
423 let row = self
424 .conn
425 .query_row(
426 "SELECT * FROM Services WHERE stable_id = ?1",
427 params![service_id],
428 service_from_row,
429 )
430 .optional()?;
431 Ok(row)
432 }
433
434 fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>> {
435 let mut stmt = self
436 .conn
437 .prepare("SELECT * FROM Services WHERE asset_id = ?1 ORDER BY port")?;
438 let rows = stmt.query_map(params![asset_id], service_from_row)?;
439 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
440 }
441
442 fn list_all_services(&mut self) -> Result<Vec<Service>> {
443 let mut stmt = self
444 .conn
445 .prepare("SELECT * FROM Services ORDER BY asset_id, port")?;
446 let rows = stmt.query_map([], service_from_row)?;
447 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
448 }
449
450 fn upsert_finding(&mut self, finding: &Finding) -> Result<()> {
451 self.conn.execute(
452 "INSERT OR REPLACE INTO Findings (stable_id, asset_id, service_id, title, severity, tool, target_url, detail, detail_path, reference, cves, endpoints, technologies, first_seen, last_seen, status, exploitability, metadata)
453 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
454 params![
455 finding.stable_id,
456 finding.asset_id,
457 finding.service_id,
458 finding.title,
459 finding.severity,
460 finding.tool,
461 finding.target_url,
462 finding.detail,
463 finding.detail_path,
464 finding.reference,
465 json_or_null(&finding.cves)?,
466 json_or_null(&finding.endpoints)?,
467 json_or_null(&finding.technologies)?,
468 finding.first_seen,
469 finding.last_seen,
470 finding.status,
471 finding.exploitability,
472 json_or_null(&finding.metadata)?
473 ],
474 )?;
475 Ok(())
476 }
477
478 fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>> {
479 let row = self
480 .conn
481 .query_row(
482 "SELECT * FROM Findings WHERE stable_id = ?1",
483 params![finding_id],
484 finding_from_row,
485 )
486 .optional()?;
487 Ok(row)
488 }
489
490 fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>> {
491 let mut stmt = self
492 .conn
493 .prepare("SELECT * FROM Findings WHERE asset_id = ?1 ORDER BY severity, title")?;
494 let rows = stmt.query_map(params![asset_id], finding_from_row)?;
495 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
496 }
497
498 fn list_all_findings(&mut self) -> Result<Vec<Finding>> {
499 let mut stmt = self
500 .conn
501 .prepare("SELECT * FROM Findings ORDER BY asset_id, title")?;
502 let rows = stmt.query_map([], finding_from_row)?;
503 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
504 }
505
506 fn add_observation(&mut self, observation: &Observation) -> Result<()> {
507 if self.get_observation(&observation.stable_id)?.is_some() {
508 return Ok(());
509 }
510 self.conn.execute(
511 "INSERT OR REPLACE INTO Observations (stable_id, run_id, subject_type, subject_id, kind, before, after, detail, observed_at)
512 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
513 params![
514 observation.stable_id,
515 observation.run_id,
516 observation.subject_type,
517 observation.subject_id,
518 observation.kind,
519 json_or_null(&observation.before)?,
520 json_or_null(&observation.after)?,
521 observation.detail,
522 observation.observed_at
523 ],
524 )?;
525 Ok(())
526 }
527
528 fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>> {
529 let row = self
530 .conn
531 .query_row(
532 "SELECT * FROM Observations WHERE stable_id = ?1",
533 params![obs_id],
534 observation_from_row,
535 )
536 .optional()?;
537 Ok(row)
538 }
539
540 fn list_observations_for(
541 &mut self,
542 subject_type: &str,
543 subject_id: &str,
544 ) -> Result<Vec<Observation>> {
545 let mut stmt = self.conn.prepare(
546 "SELECT * FROM Observations WHERE subject_type = ?1 AND subject_id = ?2 ORDER BY observed_at",
547 )?;
548 let rows = stmt.query_map(params![subject_type, subject_id], observation_from_row)?;
549 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
550 }
551
552 fn list_all_observations(&mut self) -> Result<Vec<Observation>> {
553 let mut stmt = self
554 .conn
555 .prepare("SELECT * FROM Observations ORDER BY observed_at")?;
556 let rows = stmt.query_map([], observation_from_row)?;
557 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
558 }
559
560 fn upsert_exploit_execution(&mut self, execution: &ExploitExecution) -> Result<()> {
561 self.conn.execute(
562 "INSERT OR REPLACE INTO ExploitExecutions (execution_id, run_id, asset_id, vulnerability_id, module_type, module, host, payload, status, start_time, finish_time, job_id, session_id, error, loot, options, selected, created_at)
563 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
564 params![
565 execution.execution_id,
566 execution.run_id,
567 execution.asset_id,
568 execution.vulnerability_id,
569 execution.module_type,
570 execution.module,
571 execution.host,
572 execution.payload,
573 execution.status,
574 execution.start_time,
575 execution.finish_time,
576 execution.job_id,
577 execution.session_id,
578 execution.error,
579 json_or_null(&execution.loot)?,
580 json_or_null(&execution.options)?,
581 execution.selected,
582 execution.created_at
583 ],
584 )?;
585 Ok(())
586 }
587
588 fn get_exploit_execution(&mut self, execution_id: &str) -> Result<Option<ExploitExecution>> {
589 let row = self
590 .conn
591 .query_row(
592 "SELECT * FROM ExploitExecutions WHERE execution_id = ?1",
593 params![execution_id],
594 exploit_execution_from_row,
595 )
596 .optional()?;
597 Ok(row)
598 }
599
600 fn list_exploit_executions(&mut self) -> Result<Vec<ExploitExecution>> {
601 let mut stmt = self
602 .conn
603 .prepare("SELECT * FROM ExploitExecutions ORDER BY created_at DESC")?;
604 let rows = stmt.query_map([], exploit_execution_from_row)?;
605 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
606 }
607
608 fn list_exploit_executions_for_run(&mut self, run_id: &str) -> Result<Vec<ExploitExecution>> {
609 let mut stmt = self.conn.prepare(
610 "SELECT * FROM ExploitExecutions WHERE run_id = ?1 ORDER BY created_at DESC",
611 )?;
612 let rows = stmt.query_map(params![run_id], exploit_execution_from_row)?;
613 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
614 }
615}