1use std::path::Path;
2
3use rusqlite::{params, Connection, OptionalExtension, Row};
4use serde_json::Value;
5
6use crate::db::models::{Asset, 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 INDEX IF NOT EXISTS idx_services_asset ON Services(asset_id);
105 CREATE INDEX IF NOT EXISTS idx_findings_asset ON Findings(asset_id);
106 CREATE INDEX IF NOT EXISTS idx_obs_subject ON Observations(subject_type, subject_id);",
107 )?;
108 Ok(())
109 }
110}
111
112fn json_or_null<T: serde::Serialize>(value: &T) -> rusqlite::Result<String> {
113 serde_json::to_string(value).map_err(|e| rusqlite::Error::ToSqlConversionFailure(e.into()))
114}
115
116fn parse_json(raw: Option<String>) -> Value {
117 raw.and_then(|s| serde_json::from_str(&s).ok())
118 .unwrap_or(Value::Null)
119}
120
121fn parse_vec(raw: Option<String>) -> Vec<String> {
122 match parse_json(raw) {
123 Value::Array(items) => items
124 .into_iter()
125 .filter_map(|v| v.as_str().map(str::to_string))
126 .collect(),
127 _ => Vec::new(),
128 }
129}
130
131fn parse_vec_value(raw: Option<String>) -> Vec<Value> {
132 match parse_json(raw) {
133 Value::Array(items) => items,
134 _ => Vec::new(),
135 }
136}
137
138fn asset_from_row(row: &Row) -> rusqlite::Result<Asset> {
143 Ok(Asset {
144 stable_id: row.get("stable_id")?,
145 asset_type: row.get("asset_type")?,
146 name: row.get("name")?,
147 ip: row.get("ip")?,
148 fqdn: row.get("fqdn")?,
149 reverse_dns: row.get("reverse_dns")?,
150 first_seen: row.get("first_seen")?,
151 last_seen: row.get("last_seen")?,
152 status: row.get("status")?,
153 metadata: parse_json(row.get("metadata")?),
154 })
155}
156
157fn service_from_row(row: &Row) -> rusqlite::Result<Service> {
158 Ok(Service {
159 stable_id: row.get("stable_id")?,
160 asset_id: row.get("asset_id")?,
161 port: row.get("port")?,
162 protocol: row.get("protocol")?,
163 service_name: row.get("service_name")?,
164 product: row.get("product")?,
165 version: row.get("version")?,
166 version_str: row.get("version_str")?,
167 banner: row.get("banner")?,
168 technologies: parse_vec(row.get("technologies")?),
169 cpes: parse_vec(row.get("cpes")?),
170 first_seen: row.get("first_seen")?,
171 last_seen: row.get("last_seen")?,
172 status: row.get("status")?,
173 metadata: parse_json(row.get("metadata")?),
174 })
175}
176
177fn finding_from_row(row: &Row) -> rusqlite::Result<Finding> {
178 Ok(Finding {
179 stable_id: row.get("stable_id")?,
180 asset_id: row.get("asset_id")?,
181 service_id: row.get("service_id")?,
182 title: row.get("title")?,
183 severity: row.get("severity")?,
184 tool: row.get("tool")?,
185 target_url: row.get("target_url")?,
186 detail: row.get("detail")?,
187 detail_path: row.get("detail_path")?,
188 reference: row.get("reference")?,
189 cves: parse_vec(row.get("cves")?),
190 endpoints: parse_vec(row.get("endpoints")?),
191 technologies: parse_vec(row.get("technologies")?),
192 first_seen: row.get("first_seen")?,
193 last_seen: row.get("last_seen")?,
194 status: row.get("status")?,
195 exploitability: row.get("exploitability")?,
196 metadata: parse_json(row.get("metadata")?),
197 })
198}
199
200fn observation_from_row(row: &Row) -> rusqlite::Result<Observation> {
201 Ok(Observation {
202 stable_id: row.get("stable_id")?,
203 run_id: row.get("run_id")?,
204 subject_type: row.get("subject_type")?,
205 subject_id: row.get("subject_id")?,
206 kind: row.get("kind")?,
207 before: parse_json(row.get("before")?),
208 after: parse_json(row.get("after")?),
209 detail: row.get("detail")?,
210 observed_at: row.get("observed_at")?,
211 })
212}
213
214fn scan_run_from_row(row: &Row) -> rusqlite::Result<ScanRun> {
215 Ok(ScanRun {
216 stable_id: row.get("stable_id")?,
217 target: row.get("target")?,
218 folder: row.get("folder")?,
219 started_at: row.get("started_at")?,
220 finished_at: row.get("finished_at")?,
221 tools: parse_vec(row.get("tools")?),
222 artifacts: parse_vec_value(row.get("artifacts")?),
223 content_hash: row.get("content_hash")?,
224 imported_at: row.get("imported_at")?,
225 stats: parse_json(row.get("stats")?),
226 })
227}
228
229impl Storage for SqliteStorage {
230 fn kind(&self) -> &str {
231 "sqlite"
232 }
233
234 fn ping(&mut self) -> Result<bool> {
235 let value: String = self
236 .conn
237 .query_row(
238 "SELECT name FROM sqlite_master WHERE type='table' AND name=?1 LIMIT 1",
239 params![collections::ASSETS],
240 |row| row.get(0),
241 )
242 .optional()?
243 .unwrap_or_default();
244 Ok(value == collections::ASSETS)
245 }
246
247 fn ensure_schema(&mut self) -> Result<()> {
248 self.create_tables()
249 }
250
251 fn upsert_scan_run(&mut self, run: &ScanRun) -> Result<()> {
252 self.conn.execute(
253 "INSERT OR REPLACE INTO ScanRuns (stable_id, target, folder, started_at, finished_at, tools, artifacts, content_hash, imported_at, stats)
254 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
255 params![
256 run.stable_id,
257 run.target,
258 run.folder,
259 run.started_at,
260 run.finished_at,
261 json_or_null(&run.tools)?,
262 json_or_null(&run.artifacts)?,
263 run.content_hash,
264 run.imported_at,
265 json_or_null(&run.stats)?
266 ],
267 )?;
268 Ok(())
269 }
270
271 fn get_scan_run(&mut self, run_id: &str) -> Result<Option<ScanRun>> {
272 let row = self
273 .conn
274 .query_row(
275 "SELECT * FROM ScanRuns WHERE stable_id = ?1",
276 params![run_id],
277 scan_run_from_row,
278 )
279 .optional()?;
280 Ok(row)
281 }
282
283 fn list_scan_runs(&mut self) -> Result<Vec<ScanRun>> {
284 let mut stmt = self
285 .conn
286 .prepare("SELECT * FROM ScanRuns ORDER BY imported_at DESC")?;
287 let rows = stmt.query_map([], scan_run_from_row)?;
288 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
289 }
290
291 fn upsert_asset(&mut self, asset: &Asset) -> Result<()> {
292 self.conn.execute(
293 "INSERT OR REPLACE INTO Assets (stable_id, asset_type, name, ip, fqdn, reverse_dns, first_seen, last_seen, status, metadata)
294 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
295 params![
296 asset.stable_id,
297 asset.asset_type,
298 asset.name,
299 asset.ip,
300 asset.fqdn,
301 asset.reverse_dns,
302 asset.first_seen,
303 asset.last_seen,
304 asset.status,
305 json_or_null(&asset.metadata)?
306 ],
307 )?;
308 Ok(())
309 }
310
311 fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>> {
312 let row = self
313 .conn
314 .query_row(
315 "SELECT * FROM Assets WHERE stable_id = ?1",
316 params![asset_id],
317 asset_from_row,
318 )
319 .optional()?;
320 Ok(row)
321 }
322
323 fn list_assets(&mut self) -> Result<Vec<Asset>> {
324 let mut stmt = self.conn.prepare("SELECT * FROM Assets ORDER BY name")?;
325 let rows = stmt.query_map([], asset_from_row)?;
326 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
327 }
328
329 fn upsert_service(&mut self, service: &Service) -> Result<()> {
330 self.conn.execute(
331 "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)
332 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)",
333 params![
334 service.stable_id,
335 service.asset_id,
336 service.port,
337 service.protocol,
338 service.service_name,
339 service.product,
340 service.version,
341 service.version_str,
342 service.banner,
343 json_or_null(&service.technologies)?,
344 json_or_null(&service.cpes)?,
345 service.first_seen,
346 service.last_seen,
347 service.status,
348 json_or_null(&service.metadata)?
349 ],
350 )?;
351 Ok(())
352 }
353
354 fn get_service(&mut self, service_id: &str) -> Result<Option<Service>> {
355 let row = self
356 .conn
357 .query_row(
358 "SELECT * FROM Services WHERE stable_id = ?1",
359 params![service_id],
360 service_from_row,
361 )
362 .optional()?;
363 Ok(row)
364 }
365
366 fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>> {
367 let mut stmt = self
368 .conn
369 .prepare("SELECT * FROM Services WHERE asset_id = ?1 ORDER BY port")?;
370 let rows = stmt.query_map(params![asset_id], service_from_row)?;
371 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
372 }
373
374 fn list_all_services(&mut self) -> Result<Vec<Service>> {
375 let mut stmt = self
376 .conn
377 .prepare("SELECT * FROM Services ORDER BY asset_id, port")?;
378 let rows = stmt.query_map([], service_from_row)?;
379 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
380 }
381
382 fn upsert_finding(&mut self, finding: &Finding) -> Result<()> {
383 self.conn.execute(
384 "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)
385 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
386 params![
387 finding.stable_id,
388 finding.asset_id,
389 finding.service_id,
390 finding.title,
391 finding.severity,
392 finding.tool,
393 finding.target_url,
394 finding.detail,
395 finding.detail_path,
396 finding.reference,
397 json_or_null(&finding.cves)?,
398 json_or_null(&finding.endpoints)?,
399 json_or_null(&finding.technologies)?,
400 finding.first_seen,
401 finding.last_seen,
402 finding.status,
403 finding.exploitability,
404 json_or_null(&finding.metadata)?
405 ],
406 )?;
407 Ok(())
408 }
409
410 fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>> {
411 let row = self
412 .conn
413 .query_row(
414 "SELECT * FROM Findings WHERE stable_id = ?1",
415 params![finding_id],
416 finding_from_row,
417 )
418 .optional()?;
419 Ok(row)
420 }
421
422 fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>> {
423 let mut stmt = self
424 .conn
425 .prepare("SELECT * FROM Findings WHERE asset_id = ?1 ORDER BY severity, title")?;
426 let rows = stmt.query_map(params![asset_id], finding_from_row)?;
427 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
428 }
429
430 fn list_all_findings(&mut self) -> Result<Vec<Finding>> {
431 let mut stmt = self
432 .conn
433 .prepare("SELECT * FROM Findings ORDER BY asset_id, title")?;
434 let rows = stmt.query_map([], finding_from_row)?;
435 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
436 }
437
438 fn add_observation(&mut self, observation: &Observation) -> Result<()> {
439 if self.get_observation(&observation.stable_id)?.is_some() {
440 return Ok(());
441 }
442 self.conn.execute(
443 "INSERT OR REPLACE INTO Observations (stable_id, run_id, subject_type, subject_id, kind, before, after, detail, observed_at)
444 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
445 params![
446 observation.stable_id,
447 observation.run_id,
448 observation.subject_type,
449 observation.subject_id,
450 observation.kind,
451 json_or_null(&observation.before)?,
452 json_or_null(&observation.after)?,
453 observation.detail,
454 observation.observed_at
455 ],
456 )?;
457 Ok(())
458 }
459
460 fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>> {
461 let row = self
462 .conn
463 .query_row(
464 "SELECT * FROM Observations WHERE stable_id = ?1",
465 params![obs_id],
466 observation_from_row,
467 )
468 .optional()?;
469 Ok(row)
470 }
471
472 fn list_observations_for(
473 &mut self,
474 subject_type: &str,
475 subject_id: &str,
476 ) -> Result<Vec<Observation>> {
477 let mut stmt = self.conn.prepare(
478 "SELECT * FROM Observations WHERE subject_type = ?1 AND subject_id = ?2 ORDER BY observed_at",
479 )?;
480 let rows = stmt.query_map(params![subject_type, subject_id], observation_from_row)?;
481 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
482 }
483
484 fn list_all_observations(&mut self) -> Result<Vec<Observation>> {
485 let mut stmt = self
486 .conn
487 .prepare("SELECT * FROM Observations ORDER BY observed_at")?;
488 let rows = stmt.query_map([], observation_from_row)?;
489 Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
490 }
491}