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 upsert_asset(&mut self, asset: &Asset) -> Result<()> {
346 self.conn.execute(
347 "INSERT OR REPLACE INTO Assets (stable_id, asset_type, name, ip, fqdn, reverse_dns, first_seen, last_seen, status, metadata)
348 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
349 params![
350 asset.stable_id,
351 asset.asset_type,
352 asset.name,
353 asset.ip,
354 asset.fqdn,
355 asset.reverse_dns,
356 asset.first_seen,
357 asset.last_seen,
358 asset.status,
359 json_or_null(&asset.metadata)?
360 ],
361 )?;
362 Ok(())
363 }
364
365 fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>> {
366 let row = self
367 .conn
368 .query_row(
369 "SELECT * FROM Assets WHERE stable_id = ?1",
370 params![asset_id],
371 asset_from_row,
372 )
373 .optional()?;
374 Ok(row)
375 }
376
377 fn list_assets(&mut self) -> Result<Vec<Asset>> {
378 let mut stmt = self.conn.prepare("SELECT * FROM Assets ORDER BY name")?;
379 let rows = stmt.query_map([], asset_from_row)?;
380 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
381 }
382
383 fn upsert_service(&mut self, service: &Service) -> Result<()> {
384 self.conn.execute(
385 "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)
386 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)",
387 params![
388 service.stable_id,
389 service.asset_id,
390 service.port,
391 service.protocol,
392 service.service_name,
393 service.product,
394 service.version,
395 service.version_str,
396 service.banner,
397 json_or_null(&service.technologies)?,
398 json_or_null(&service.cpes)?,
399 service.first_seen,
400 service.last_seen,
401 service.status,
402 json_or_null(&service.metadata)?
403 ],
404 )?;
405 Ok(())
406 }
407
408 fn get_service(&mut self, service_id: &str) -> Result<Option<Service>> {
409 let row = self
410 .conn
411 .query_row(
412 "SELECT * FROM Services WHERE stable_id = ?1",
413 params![service_id],
414 service_from_row,
415 )
416 .optional()?;
417 Ok(row)
418 }
419
420 fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>> {
421 let mut stmt = self
422 .conn
423 .prepare("SELECT * FROM Services WHERE asset_id = ?1 ORDER BY port")?;
424 let rows = stmt.query_map(params![asset_id], service_from_row)?;
425 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
426 }
427
428 fn list_all_services(&mut self) -> Result<Vec<Service>> {
429 let mut stmt = self
430 .conn
431 .prepare("SELECT * FROM Services ORDER BY asset_id, port")?;
432 let rows = stmt.query_map([], service_from_row)?;
433 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
434 }
435
436 fn upsert_finding(&mut self, finding: &Finding) -> Result<()> {
437 self.conn.execute(
438 "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)
439 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
440 params![
441 finding.stable_id,
442 finding.asset_id,
443 finding.service_id,
444 finding.title,
445 finding.severity,
446 finding.tool,
447 finding.target_url,
448 finding.detail,
449 finding.detail_path,
450 finding.reference,
451 json_or_null(&finding.cves)?,
452 json_or_null(&finding.endpoints)?,
453 json_or_null(&finding.technologies)?,
454 finding.first_seen,
455 finding.last_seen,
456 finding.status,
457 finding.exploitability,
458 json_or_null(&finding.metadata)?
459 ],
460 )?;
461 Ok(())
462 }
463
464 fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>> {
465 let row = self
466 .conn
467 .query_row(
468 "SELECT * FROM Findings WHERE stable_id = ?1",
469 params![finding_id],
470 finding_from_row,
471 )
472 .optional()?;
473 Ok(row)
474 }
475
476 fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>> {
477 let mut stmt = self
478 .conn
479 .prepare("SELECT * FROM Findings WHERE asset_id = ?1 ORDER BY severity, title")?;
480 let rows = stmt.query_map(params![asset_id], finding_from_row)?;
481 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
482 }
483
484 fn list_all_findings(&mut self) -> Result<Vec<Finding>> {
485 let mut stmt = self
486 .conn
487 .prepare("SELECT * FROM Findings ORDER BY asset_id, title")?;
488 let rows = stmt.query_map([], finding_from_row)?;
489 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
490 }
491
492 fn add_observation(&mut self, observation: &Observation) -> Result<()> {
493 if self.get_observation(&observation.stable_id)?.is_some() {
494 return Ok(());
495 }
496 self.conn.execute(
497 "INSERT OR REPLACE INTO Observations (stable_id, run_id, subject_type, subject_id, kind, before, after, detail, observed_at)
498 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
499 params![
500 observation.stable_id,
501 observation.run_id,
502 observation.subject_type,
503 observation.subject_id,
504 observation.kind,
505 json_or_null(&observation.before)?,
506 json_or_null(&observation.after)?,
507 observation.detail,
508 observation.observed_at
509 ],
510 )?;
511 Ok(())
512 }
513
514 fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>> {
515 let row = self
516 .conn
517 .query_row(
518 "SELECT * FROM Observations WHERE stable_id = ?1",
519 params![obs_id],
520 observation_from_row,
521 )
522 .optional()?;
523 Ok(row)
524 }
525
526 fn list_observations_for(
527 &mut self,
528 subject_type: &str,
529 subject_id: &str,
530 ) -> Result<Vec<Observation>> {
531 let mut stmt = self.conn.prepare(
532 "SELECT * FROM Observations WHERE subject_type = ?1 AND subject_id = ?2 ORDER BY observed_at",
533 )?;
534 let rows = stmt.query_map(params![subject_type, subject_id], observation_from_row)?;
535 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
536 }
537
538 fn list_all_observations(&mut self) -> Result<Vec<Observation>> {
539 let mut stmt = self
540 .conn
541 .prepare("SELECT * FROM Observations ORDER BY observed_at")?;
542 let rows = stmt.query_map([], observation_from_row)?;
543 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
544 }
545
546 fn upsert_exploit_execution(&mut self, execution: &ExploitExecution) -> Result<()> {
547 self.conn.execute(
548 "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)
549 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
550 params![
551 execution.execution_id,
552 execution.run_id,
553 execution.asset_id,
554 execution.vulnerability_id,
555 execution.module_type,
556 execution.module,
557 execution.host,
558 execution.payload,
559 execution.status,
560 execution.start_time,
561 execution.finish_time,
562 execution.job_id,
563 execution.session_id,
564 execution.error,
565 json_or_null(&execution.loot)?,
566 json_or_null(&execution.options)?,
567 execution.selected,
568 execution.created_at
569 ],
570 )?;
571 Ok(())
572 }
573
574 fn get_exploit_execution(&mut self, execution_id: &str) -> Result<Option<ExploitExecution>> {
575 let row = self
576 .conn
577 .query_row(
578 "SELECT * FROM ExploitExecutions WHERE execution_id = ?1",
579 params![execution_id],
580 exploit_execution_from_row,
581 )
582 .optional()?;
583 Ok(row)
584 }
585
586 fn list_exploit_executions(&mut self) -> Result<Vec<ExploitExecution>> {
587 let mut stmt = self
588 .conn
589 .prepare("SELECT * FROM ExploitExecutions ORDER BY created_at DESC")?;
590 let rows = stmt.query_map([], exploit_execution_from_row)?;
591 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
592 }
593
594 fn list_exploit_executions_for_run(&mut self, run_id: &str) -> Result<Vec<ExploitExecution>> {
595 let mut stmt = self.conn.prepare(
596 "SELECT * FROM ExploitExecutions WHERE run_id = ?1 ORDER BY created_at DESC",
597 )?;
598 let rows = stmt.query_map(params![run_id], exploit_execution_from_row)?;
599 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
600 }
601}