Skip to main content

ploit_malper/db/
mod.rs

1pub mod content;
2pub mod models;
3pub mod pocketbase;
4pub mod schema;
5pub mod sqlite;
6
7use crate::config::DatabaseConfig;
8use crate::error::Result;
9use models::{Asset, ExploitExecution, Finding, Observation, ScanRun, Service};
10
11/// Collection names used by every backend. Kept intentionally small: bulk
12/// content (full reports, long finding evidence) is stored as files in the
13/// content store, so the database only holds metadata and short fields.
14pub mod collections {
15    pub const ASSETS: &str = "Assets";
16    pub const SERVICES: &str = "Services";
17    pub const FINDINGS: &str = "Findings";
18    pub const OBSERVATIONS: &str = "Observations";
19    pub const SCAN_RUNS: &str = "ScanRuns";
20    pub const EXPLOIT_EXECUTIONS: &str = "ExploitExecutions";
21
22    pub const ALL: [&str; 6] = [
23        ASSETS,
24        SERVICES,
25        FINDINGS,
26        OBSERVATIONS,
27        SCAN_RUNS,
28        EXPLOIT_EXECUTIONS,
29    ];
30}
31
32/// Backend-agnostic persistence interface used by the importer.
33///
34/// All write operations are upserts keyed by a stable, content-derived id so
35/// that re-running an import over the same scan folder is idempotent. History
36/// is preserved through the `Observations` collection rather than by
37/// overwriting prior rows.
38pub trait Storage: Send {
39    fn kind(&self) -> &str;
40
41    /// Verify connectivity (PocketBase health / SQLite open + query).
42    fn ping(&mut self) -> Result<bool>;
43
44    /// Create all required collections/tables if they do not exist.
45    fn ensure_schema(&mut self) -> Result<()>;
46
47    // --- Scan runs ------------------------------------------------------
48
49    fn upsert_scan_run(&mut self, run: &ScanRun) -> Result<()>;
50    fn get_scan_run(&mut self, run_id: &str) -> Result<Option<ScanRun>>;
51    fn list_scan_runs(&mut self) -> Result<Vec<ScanRun>>;
52    /// Delete a scan run and the records owned by it (observations and
53    /// exploit executions). Assets, services and findings are shared across
54    /// runs and are left untouched.
55    fn delete_scan_run(&mut self, run_id: &str) -> Result<()>;
56
57    // --- Assets ---------------------------------------------------------
58
59    fn upsert_asset(&mut self, asset: &Asset) -> Result<()>;
60    fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>>;
61    fn list_assets(&mut self) -> Result<Vec<Asset>>;
62
63    // --- Services -------------------------------------------------------
64
65    fn upsert_service(&mut self, service: &Service) -> Result<()>;
66    fn get_service(&mut self, service_id: &str) -> Result<Option<Service>>;
67    fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>>;
68    fn list_all_services(&mut self) -> Result<Vec<Service>>;
69
70    // --- Findings -------------------------------------------------------
71
72    fn upsert_finding(&mut self, finding: &Finding) -> Result<()>;
73    fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>>;
74    fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>>;
75    fn list_all_findings(&mut self) -> Result<Vec<Finding>>;
76
77    // --- Observations ---------------------------------------------------
78
79    fn add_observation(&mut self, observation: &Observation) -> Result<()>;
80    fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>>;
81    fn list_observations_for(
82        &mut self,
83        subject_type: &str,
84        subject_id: &str,
85    ) -> Result<Vec<Observation>>;
86    fn list_all_observations(&mut self) -> Result<Vec<Observation>>;
87
88    // --- Exploit executions ---------------------------------------------
89
90    fn upsert_exploit_execution(&mut self, execution: &ExploitExecution) -> Result<()>;
91    fn get_exploit_execution(&mut self, execution_id: &str) -> Result<Option<ExploitExecution>>;
92    fn list_exploit_executions(&mut self) -> Result<Vec<ExploitExecution>>;
93    fn list_exploit_executions_for_run(&mut self, run_id: &str) -> Result<Vec<ExploitExecution>>;
94}
95
96/// Open a storage backend based on the persisted configuration.
97pub fn open_storage(config: &DatabaseConfig) -> Result<Box<dyn Storage>> {
98    if config.backend == "sqlite" {
99        Ok(Box::new(sqlite::SqliteStorage::open(&config.sqlite_path)?))
100    } else {
101        Ok(Box::new(pocketbase::PocketBaseStorage::new(config)?))
102    }
103}
104
105/// Open a storage backend, overriding the configured backend.
106pub fn open_storage_with(
107    config: &DatabaseConfig,
108    backend: &str,
109    pocketbase_url: Option<&str>,
110    sqlite_path: Option<&str>,
111) -> Result<Box<dyn Storage>> {
112    if backend == "sqlite" {
113        let path = sqlite_path
114            .map(str::to_string)
115            .unwrap_or_else(|| config.sqlite_path.clone());
116        Ok(Box::new(sqlite::SqliteStorage::open(&path)?))
117    } else {
118        let mut override_config = config.clone();
119        if let Some(url) = pocketbase_url {
120            override_config.pocketbase_url = url.to_string();
121        }
122        Ok(Box::new(pocketbase::PocketBaseStorage::new(
123            &override_config,
124        )?))
125    }
126}