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
11pub 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
32pub trait Storage: Send {
39 fn kind(&self) -> &str;
40
41 fn ping(&mut self) -> Result<bool>;
43
44 fn ensure_schema(&mut self) -> Result<()>;
46
47 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 fn delete_scan_run(&mut self, run_id: &str) -> Result<()>;
56
57 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 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 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 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 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
96pub 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
105pub 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}