ploit_malper/db/
sqlite.rs

1use std::path::Path;
2
3use rusqlite::{params, Connection, OptionalExtension, Row};
4use serde_json::Value;
5
6use crate::db::models::{Asset, Finding, Observation, Relationship, Report, ScanRun, Service};
7use crate::db::{collections, Storage};
8use crate::error::Result;
9
10/// SQLite storage backend. Each collection maps to a table; JSON fields are
11/// stored as TEXT columns and (de)serialized on write/read.
12pub 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                reference       TEXT,
71                cves            TEXT,
72                endpoints       TEXT,
73                technologies    TEXT,
74                first_seen      TEXT,
75                last_seen       TEXT,
76                status          TEXT,
77                exploitability  TEXT,
78                metadata        TEXT
79            );
80            CREATE TABLE IF NOT EXISTS Observations (
81                stable_id    TEXT PRIMARY KEY,
82                run_id       TEXT NOT NULL,
83                subject_type TEXT NOT NULL,
84                subject_id   TEXT NOT NULL,
85                kind         TEXT NOT NULL,
86                before       TEXT,
87                after        TEXT,
88                detail       TEXT,
89                observed_at  TEXT
90            );
91            CREATE TABLE IF NOT EXISTS Relationships (
92                stable_id    TEXT PRIMARY KEY,
93                run_id       TEXT,
94                subject_type TEXT NOT NULL,
95                subject_id   TEXT NOT NULL,
96                object_type  TEXT NOT NULL,
97                object_id    TEXT NOT NULL,
98                kind         TEXT NOT NULL,
99                detail       TEXT
100            );
101            CREATE TABLE IF NOT EXISTS ScanRuns (
102                stable_id    TEXT PRIMARY KEY,
103                target       TEXT NOT NULL,
104                folder       TEXT,
105                started_at   TEXT,
106                finished_at  TEXT,
107                tools        TEXT,
108                artifacts    TEXT,
109                content_hash TEXT,
110                imported_at  TEXT,
111                stats        TEXT
112            );
113            CREATE TABLE IF NOT EXISTS Reports (
114                stable_id    TEXT PRIMARY KEY,
115                run_id       TEXT NOT NULL,
116                tool         TEXT NOT NULL,
117                format       TEXT NOT NULL,
118                title        TEXT,
119                source_path  TEXT,
120                content      TEXT,
121                content_hash TEXT,
122                imported_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_rel_subject ON Relationships(subject_type, subject_id);
128            CREATE INDEX IF NOT EXISTS idx_reports_run ON Reports(run_id);",
129        )?;
130        Ok(())
131    }
132}
133
134fn json_or_null<T: serde::Serialize>(value: &T) -> rusqlite::Result<String> {
135    serde_json::to_string(value).map_err(|e| rusqlite::Error::ToSqlConversionFailure(e.into()))
136}
137
138fn parse_json(raw: Option<String>) -> Value {
139    raw.and_then(|s| serde_json::from_str(&s).ok())
140        .unwrap_or(Value::Null)
141}
142
143fn parse_vec(raw: Option<String>) -> Vec<String> {
144    match parse_json(raw) {
145        Value::Array(items) => items
146            .into_iter()
147            .filter_map(|v| v.as_str().map(str::to_string))
148            .collect(),
149        _ => Vec::new(),
150    }
151}
152
153fn parse_vec_value(raw: Option<String>) -> Vec<Value> {
154    match parse_json(raw) {
155        Value::Array(items) => items,
156        _ => Vec::new(),
157    }
158}
159
160// ---------------------------------------------------------------------------
161// Row mapping
162// ---------------------------------------------------------------------------
163
164fn asset_from_row(row: &Row) -> rusqlite::Result<Asset> {
165    Ok(Asset {
166        stable_id: row.get("stable_id")?,
167        asset_type: row.get("asset_type")?,
168        name: row.get("name")?,
169        ip: row.get("ip")?,
170        fqdn: row.get("fqdn")?,
171        reverse_dns: row.get("reverse_dns")?,
172        first_seen: row.get("first_seen")?,
173        last_seen: row.get("last_seen")?,
174        status: row.get("status")?,
175        metadata: parse_json(row.get("metadata")?),
176    })
177}
178
179fn service_from_row(row: &Row) -> rusqlite::Result<Service> {
180    Ok(Service {
181        stable_id: row.get("stable_id")?,
182        asset_id: row.get("asset_id")?,
183        port: row.get("port")?,
184        protocol: row.get("protocol")?,
185        service_name: row.get("service_name")?,
186        product: row.get("product")?,
187        version: row.get("version")?,
188        version_str: row.get("version_str")?,
189        banner: row.get("banner")?,
190        technologies: parse_vec(row.get("technologies")?),
191        cpes: parse_vec(row.get("cpes")?),
192        first_seen: row.get("first_seen")?,
193        last_seen: row.get("last_seen")?,
194        status: row.get("status")?,
195        metadata: parse_json(row.get("metadata")?),
196    })
197}
198
199fn finding_from_row(row: &Row) -> rusqlite::Result<Finding> {
200    Ok(Finding {
201        stable_id: row.get("stable_id")?,
202        asset_id: row.get("asset_id")?,
203        service_id: row.get("service_id")?,
204        title: row.get("title")?,
205        severity: row.get("severity")?,
206        tool: row.get("tool")?,
207        target_url: row.get("target_url")?,
208        detail: row.get("detail")?,
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 relationship_from_row(row: &Row) -> rusqlite::Result<Relationship> {
236    Ok(Relationship {
237        stable_id: row.get("stable_id")?,
238        run_id: row.get("run_id")?,
239        subject_type: row.get("subject_type")?,
240        subject_id: row.get("subject_id")?,
241        object_type: row.get("object_type")?,
242        object_id: row.get("object_id")?,
243        kind: row.get("kind")?,
244        detail: row.get("detail")?,
245    })
246}
247
248fn scan_run_from_row(row: &Row) -> rusqlite::Result<ScanRun> {
249    Ok(ScanRun {
250        stable_id: row.get("stable_id")?,
251        target: row.get("target")?,
252        folder: row.get("folder")?,
253        started_at: row.get("started_at")?,
254        finished_at: row.get("finished_at")?,
255        tools: parse_vec(row.get("tools")?),
256        artifacts: parse_vec_value(row.get("artifacts")?),
257        content_hash: row.get("content_hash")?,
258        imported_at: row.get("imported_at")?,
259        stats: parse_json(row.get("stats")?),
260    })
261}
262
263fn report_from_row(row: &Row) -> rusqlite::Result<Report> {
264    Ok(Report {
265        stable_id: row.get("stable_id")?,
266        run_id: row.get("run_id")?,
267        tool: row.get("tool")?,
268        format: row.get("format")?,
269        title: row.get("title")?,
270        source_path: row.get("source_path")?,
271        content: row.get("content")?,
272        content_hash: row.get("content_hash")?,
273        imported_at: row.get("imported_at")?,
274    })
275}
276
277impl Storage for SqliteStorage {
278    fn kind(&self) -> &str {
279        "sqlite"
280    }
281
282    fn ping(&mut self) -> Result<bool> {
283        let value: String = self
284            .conn
285            .query_row(
286                "SELECT name FROM sqlite_master WHERE type='table' AND name=?1 LIMIT 1",
287                params![collections::ASSETS],
288                |row| row.get(0),
289            )
290            .optional()?
291            .unwrap_or_default();
292        Ok(value == collections::ASSETS)
293    }
294
295    fn ensure_schema(&mut self) -> Result<()> {
296        self.create_tables()
297    }
298
299    fn upsert_scan_run(&mut self, run: &ScanRun) -> Result<()> {
300        self.conn.execute(
301            "INSERT OR REPLACE INTO ScanRuns (stable_id, target, folder, started_at, finished_at, tools, artifacts, content_hash, imported_at, stats)
302             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
303            params![
304                run.stable_id,
305                run.target,
306                run.folder,
307                run.started_at,
308                run.finished_at,
309                json_or_null(&run.tools)?,
310                json_or_null(&run.artifacts)?,
311                run.content_hash,
312                run.imported_at,
313                json_or_null(&run.stats)?
314            ],
315        )?;
316        Ok(())
317    }
318
319    fn get_scan_run(&mut self, run_id: &str) -> Result<Option<ScanRun>> {
320        let row = self
321            .conn
322            .query_row(
323                "SELECT * FROM ScanRuns WHERE stable_id = ?1",
324                params![run_id],
325                scan_run_from_row,
326            )
327            .optional()?;
328        Ok(row)
329    }
330
331    fn list_scan_runs(&mut self) -> Result<Vec<ScanRun>> {
332        let mut stmt = self
333            .conn
334            .prepare("SELECT * FROM ScanRuns ORDER BY imported_at DESC")?;
335        let rows = stmt.query_map([], scan_run_from_row)?;
336        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
337    }
338
339    fn upsert_asset(&mut self, asset: &Asset) -> Result<()> {
340        self.conn.execute(
341            "INSERT OR REPLACE INTO Assets (stable_id, asset_type, name, ip, fqdn, reverse_dns, first_seen, last_seen, status, metadata)
342             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
343            params![
344                asset.stable_id,
345                asset.asset_type,
346                asset.name,
347                asset.ip,
348                asset.fqdn,
349                asset.reverse_dns,
350                asset.first_seen,
351                asset.last_seen,
352                asset.status,
353                json_or_null(&asset.metadata)?
354            ],
355        )?;
356        Ok(())
357    }
358
359    fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>> {
360        let row = self
361            .conn
362            .query_row(
363                "SELECT * FROM Assets WHERE stable_id = ?1",
364                params![asset_id],
365                asset_from_row,
366            )
367            .optional()?;
368        Ok(row)
369    }
370
371    fn list_assets(&mut self) -> Result<Vec<Asset>> {
372        let mut stmt = self.conn.prepare("SELECT * FROM Assets ORDER BY name")?;
373        let rows = stmt.query_map([], asset_from_row)?;
374        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
375    }
376
377    fn upsert_service(&mut self, service: &Service) -> Result<()> {
378        self.conn.execute(
379            "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)
380             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)",
381            params![
382                service.stable_id,
383                service.asset_id,
384                service.port,
385                service.protocol,
386                service.service_name,
387                service.product,
388                service.version,
389                service.version_str,
390                service.banner,
391                json_or_null(&service.technologies)?,
392                json_or_null(&service.cpes)?,
393                service.first_seen,
394                service.last_seen,
395                service.status,
396                json_or_null(&service.metadata)?
397            ],
398        )?;
399        Ok(())
400    }
401
402    fn get_service(&mut self, service_id: &str) -> Result<Option<Service>> {
403        let row = self
404            .conn
405            .query_row(
406                "SELECT * FROM Services WHERE stable_id = ?1",
407                params![service_id],
408                service_from_row,
409            )
410            .optional()?;
411        Ok(row)
412    }
413
414    fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>> {
415        let mut stmt = self
416            .conn
417            .prepare("SELECT * FROM Services WHERE asset_id = ?1 ORDER BY port")?;
418        let rows = stmt.query_map(params![asset_id], service_from_row)?;
419        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
420    }
421
422    fn list_all_services(&mut self) -> Result<Vec<Service>> {
423        let mut stmt = self
424            .conn
425            .prepare("SELECT * FROM Services ORDER BY asset_id, port")?;
426        let rows = stmt.query_map([], service_from_row)?;
427        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
428    }
429
430    fn upsert_finding(&mut self, finding: &Finding) -> Result<()> {
431        self.conn.execute(
432            "INSERT OR REPLACE INTO Findings (stable_id, asset_id, service_id, title, severity, tool, target_url, detail, reference, cves, endpoints, technologies, first_seen, last_seen, status, exploitability, metadata)
433             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17)",
434            params![
435                finding.stable_id,
436                finding.asset_id,
437                finding.service_id,
438                finding.title,
439                finding.severity,
440                finding.tool,
441                finding.target_url,
442                finding.detail,
443                finding.reference,
444                json_or_null(&finding.cves)?,
445                json_or_null(&finding.endpoints)?,
446                json_or_null(&finding.technologies)?,
447                finding.first_seen,
448                finding.last_seen,
449                finding.status,
450                finding.exploitability,
451                json_or_null(&finding.metadata)?
452            ],
453        )?;
454        Ok(())
455    }
456
457    fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>> {
458        let row = self
459            .conn
460            .query_row(
461                "SELECT * FROM Findings WHERE stable_id = ?1",
462                params![finding_id],
463                finding_from_row,
464            )
465            .optional()?;
466        Ok(row)
467    }
468
469    fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>> {
470        let mut stmt = self
471            .conn
472            .prepare("SELECT * FROM Findings WHERE asset_id = ?1 ORDER BY severity, title")?;
473        let rows = stmt.query_map(params![asset_id], finding_from_row)?;
474        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
475    }
476
477    fn list_all_findings(&mut self) -> Result<Vec<Finding>> {
478        let mut stmt = self
479            .conn
480            .prepare("SELECT * FROM Findings ORDER BY asset_id, title")?;
481        let rows = stmt.query_map([], finding_from_row)?;
482        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
483    }
484
485    fn add_observation(&mut self, observation: &Observation) -> Result<()> {
486        if self.get_observation(&observation.stable_id)?.is_some() {
487            return Ok(());
488        }
489        self.conn.execute(
490            "INSERT OR REPLACE INTO Observations (stable_id, run_id, subject_type, subject_id, kind, before, after, detail, observed_at)
491             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
492            params![
493                observation.stable_id,
494                observation.run_id,
495                observation.subject_type,
496                observation.subject_id,
497                observation.kind,
498                json_or_null(&observation.before)?,
499                json_or_null(&observation.after)?,
500                observation.detail,
501                observation.observed_at
502            ],
503        )?;
504        Ok(())
505    }
506
507    fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>> {
508        let row = self
509            .conn
510            .query_row(
511                "SELECT * FROM Observations WHERE stable_id = ?1",
512                params![obs_id],
513                observation_from_row,
514            )
515            .optional()?;
516        Ok(row)
517    }
518
519    fn list_observations_for(
520        &mut self,
521        subject_type: &str,
522        subject_id: &str,
523    ) -> Result<Vec<Observation>> {
524        let mut stmt = self.conn.prepare(
525            "SELECT * FROM Observations WHERE subject_type = ?1 AND subject_id = ?2 ORDER BY observed_at",
526        )?;
527        let rows = stmt.query_map(params![subject_type, subject_id], observation_from_row)?;
528        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
529    }
530
531    fn list_all_observations(&mut self) -> Result<Vec<Observation>> {
532        let mut stmt = self
533            .conn
534            .prepare("SELECT * FROM Observations ORDER BY observed_at")?;
535        let rows = stmt.query_map([], observation_from_row)?;
536        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
537    }
538
539    fn upsert_relationship(&mut self, relationship: &Relationship) -> Result<()> {
540        self.conn.execute(
541            "INSERT OR REPLACE INTO Relationships (stable_id, run_id, subject_type, subject_id, object_type, object_id, kind, detail)
542             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
543            params![
544                relationship.stable_id,
545                relationship.run_id,
546                relationship.subject_type,
547                relationship.subject_id,
548                relationship.object_type,
549                relationship.object_id,
550                relationship.kind,
551                relationship.detail
552            ],
553        )?;
554        Ok(())
555    }
556
557    fn get_relationship(&mut self, rel_id: &str) -> Result<Option<Relationship>> {
558        let row = self
559            .conn
560            .query_row(
561                "SELECT * FROM Relationships WHERE stable_id = ?1",
562                params![rel_id],
563                relationship_from_row,
564            )
565            .optional()?;
566        Ok(row)
567    }
568
569    fn list_relationships_for(
570        &mut self,
571        subject_type: &str,
572        subject_id: &str,
573    ) -> Result<Vec<Relationship>> {
574        let mut stmt = self.conn.prepare(
575            "SELECT * FROM Relationships WHERE subject_type = ?1 AND subject_id = ?2 ORDER BY kind",
576        )?;
577        let rows = stmt.query_map(params![subject_type, subject_id], relationship_from_row)?;
578        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
579    }
580
581    fn upsert_report(&mut self, report: &Report) -> Result<()> {
582        self.conn.execute(
583            "INSERT OR REPLACE INTO Reports (stable_id, run_id, tool, format, title, source_path, content, content_hash, imported_at)
584             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
585            params![
586                report.stable_id,
587                report.run_id,
588                report.tool,
589                report.format,
590                report.title,
591                report.source_path,
592                report.content,
593                report.content_hash,
594                report.imported_at
595            ],
596        )?;
597        Ok(())
598    }
599
600    fn get_report(&mut self, report_id: &str) -> Result<Option<Report>> {
601        let row = self
602            .conn
603            .query_row(
604                "SELECT * FROM Reports WHERE stable_id = ?1",
605                params![report_id],
606                report_from_row,
607            )
608            .optional()?;
609        Ok(row)
610    }
611
612    fn list_reports_for_run(&mut self, run_id: &str) -> Result<Vec<Report>> {
613        let mut stmt = self
614            .conn
615            .prepare("SELECT * FROM Reports WHERE run_id = ?1 ORDER BY tool")?;
616        let rows = stmt.query_map(params![run_id], report_from_row)?;
617        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
618    }
619}
ploit_malper/db/sqlite.rs

ploit_malper/db/
sqlite.rs
