ploit_malper/db/
pocketbase.rs

1use std::io::Write;
2use std::process::{Command, Stdio};
3
4use serde_json::{json, Value};
5
6use crate::config::DatabaseConfig;
7use crate::db::models::{Asset, ExploitExecution, Finding, Observation, ScanRun, Service};
8use crate::db::schema::{self, FieldDef};
9use crate::db::{collections, Storage};
10use crate::error::{AppError, Result};
11
12const REQUEST_TIMEOUT_S: u64 = 60;
13
14/// PocketBase storage backend. Talks to the PocketBase REST API over `curl`
15/// (matching the rest of PloitMalper's lightweight HTTP transport) and uses
16/// `pbctl` for declarative schema management when available.
17pub struct PocketBaseStorage {
18    config: DatabaseConfig,
19    token: String,
20    /// True once the current token has been verified against the API.
21    token_validated: bool,
22}
23
24struct HttpResponse {
25    status: u16,
26    body: Value,
27}
28
29fn curl_request(
30    method: &str,
31    url: &str,
32    token: Option<&str>,
33    body: Option<&Value>,
34    query: &[(&str, &str)],
35) -> Result<HttpResponse> {
36    let mut command = Command::new("curl");
37    command
38        .arg("-sS")
39        .arg("--max-time")
40        .arg(REQUEST_TIMEOUT_S.to_string())
41        .arg("-H")
42        .arg("Content-Type: application/json")
43        .arg("-w")
44        .arg("\n%{http_code}")
45        .stdout(Stdio::piped())
46        .stderr(Stdio::piped());
47
48    if let Some(token) = token {
49        command
50            .arg("-H")
51            .arg(format!("Authorization: Bearer {}", token));
52    }
53
54    if !query.is_empty() {
55        command.arg("-G");
56        for (key, value) in query {
57            command
58                .arg("--data-urlencode")
59                .arg(format!("{}={}", key, value));
60        }
61    }
62
63    if body.is_some() {
64        if method != "GET" {
65            command.arg("-X").arg(method);
66        }
67        command.arg("--data-binary").arg("@-").stdin(Stdio::piped());
68    } else if method != "GET" {
69        command.arg("-X").arg(method);
70    }
71
72    command.arg(url);
73
74    let mut child = command
75        .spawn()
76        .map_err(|e| AppError::Network(format!("failed to spawn curl: {e}")))?;
77
78    if let Some(b) = body {
79        if let Some(mut stdin) = child.stdin.take() {
80            let payload = serde_json::to_string(b)?;
81            stdin
82                .write_all(payload.as_bytes())
83                .map_err(|e| AppError::Network(format!("failed to write request body: {e}")))?;
84        }
85    }
86
87    let output = child
88        .wait_with_output()
89        .map_err(|e| AppError::Network(format!("failed to wait for curl: {e}")))?;
90
91    let raw = String::from_utf8_lossy(&output.stdout).to_string();
92    let trimmed = raw.trim_end();
93    let (body_str, status) = match trimmed.rsplit_once('\n') {
94        Some((rest, code)) => (rest, code.trim().parse::<u16>().unwrap_or(0)),
95        None => (trimmed, 0),
96    };
97
98    let body: Value = if body_str.trim().is_empty() {
99        Value::Null
100    } else {
101        serde_json::from_str(body_str).unwrap_or_else(|_| json!({ "raw": body_str }))
102    };
103
104    Ok(HttpResponse { status, body })
105}
106
107fn get_str(value: &Value, key: &str) -> String {
108    value
109        .get(key)
110        .and_then(Value::as_str)
111        .unwrap_or("")
112        .to_string()
113}
114
115fn get_opt_str(value: &Value, key: &str) -> Option<String> {
116    value
117        .get(key)
118        .and_then(Value::as_str)
119        .filter(|s| !s.is_empty())
120        .map(str::to_string)
121}
122
123fn get_vec(value: &Value, key: &str) -> Vec<String> {
124    value
125        .get(key)
126        .and_then(Value::as_array)
127        .map(|items| {
128            items
129                .iter()
130                .filter_map(|v| v.as_str().map(str::to_string))
131                .collect()
132        })
133        .unwrap_or_default()
134}
135
136fn get_vec_value(value: &Value, key: &str) -> Vec<Value> {
137    value
138        .get(key)
139        .and_then(Value::as_array)
140        .cloned()
141        .unwrap_or_default()
142}
143
144fn get_json(value: &Value, key: &str) -> Value {
145    value.get(key).cloned().unwrap_or(Value::Null)
146}
147
148fn as_json_object(value: &Value) -> Value {
149    match value {
150        Value::Null => Value::Object(Default::default()),
151        other => other.clone(),
152    }
153}
154
155// ---------------------------------------------------------------------------
156// Record <-> model conversion
157// ---------------------------------------------------------------------------
158
159fn asset_to_record(asset: &Asset) -> Value {
160    let mut record = serde_json::Map::new();
161    record.insert("stable_id".into(), json!(asset.stable_id));
162    record.insert("asset_type".into(), json!(asset.asset_type));
163    record.insert("name".into(), json!(asset.name));
164    if let Some(v) = &asset.ip {
165        record.insert("ip".into(), json!(v));
166    }
167    if let Some(v) = &asset.fqdn {
168        record.insert("fqdn".into(), json!(v));
169    }
170    if let Some(v) = &asset.reverse_dns {
171        record.insert("reverse_dns".into(), json!(v));
172    }
173    record.insert("first_seen".into(), json!(asset.first_seen));
174    record.insert("last_seen".into(), json!(asset.last_seen));
175    record.insert("status".into(), json!(asset.status));
176    record.insert("metadata".into(), as_json_object(&asset.metadata));
177    Value::Object(record)
178}
179
180fn record_to_asset(value: &Value) -> Asset {
181    Asset {
182        stable_id: get_str(value, "stable_id"),
183        asset_type: get_str(value, "asset_type"),
184        name: get_str(value, "name"),
185        ip: get_opt_str(value, "ip"),
186        fqdn: get_opt_str(value, "fqdn"),
187        reverse_dns: get_opt_str(value, "reverse_dns"),
188        first_seen: get_str(value, "first_seen"),
189        last_seen: get_str(value, "last_seen"),
190        status: get_str(value, "status"),
191        metadata: get_json(value, "metadata"),
192    }
193}
194
195fn service_to_record(service: &Service) -> Value {
196    let mut record = serde_json::Map::new();
197    record.insert("stable_id".into(), json!(service.stable_id));
198    record.insert("asset_id".into(), json!(service.asset_id));
199    record.insert("port".into(), json!(service.port));
200    record.insert("protocol".into(), json!(service.protocol));
201    record.insert("service_name".into(), json!(service.service_name));
202    if let Some(v) = &service.product {
203        record.insert("product".into(), json!(v));
204    }
205    if let Some(v) = &service.version {
206        record.insert("version".into(), json!(v));
207    }
208    if let Some(v) = &service.version_str {
209        record.insert("version_str".into(), json!(v));
210    }
211    if let Some(v) = &service.banner {
212        record.insert("banner".into(), json!(v));
213    }
214    record.insert("technologies".into(), json!(service.technologies));
215    record.insert("cpes".into(), json!(service.cpes));
216    record.insert("first_seen".into(), json!(service.first_seen));
217    record.insert("last_seen".into(), json!(service.last_seen));
218    record.insert("status".into(), json!(service.status));
219    record.insert("metadata".into(), as_json_object(&service.metadata));
220    Value::Object(record)
221}
222
223fn record_to_service(value: &Value) -> Service {
224    Service {
225        stable_id: get_str(value, "stable_id"),
226        asset_id: get_str(value, "asset_id"),
227        port: value
228            .get("port")
229            .and_then(Value::as_u64)
230            .and_then(|p| u16::try_from(p).ok())
231            .unwrap_or(0),
232        protocol: get_str(value, "protocol"),
233        service_name: get_str(value, "service_name"),
234        product: get_opt_str(value, "product"),
235        version: get_opt_str(value, "version"),
236        version_str: get_opt_str(value, "version_str"),
237        banner: get_opt_str(value, "banner"),
238        technologies: get_vec(value, "technologies"),
239        cpes: get_vec(value, "cpes"),
240        first_seen: get_str(value, "first_seen"),
241        last_seen: get_str(value, "last_seen"),
242        status: get_str(value, "status"),
243        metadata: get_json(value, "metadata"),
244    }
245}
246
247fn finding_to_record(finding: &Finding) -> Value {
248    let mut record = serde_json::Map::new();
249    record.insert("stable_id".into(), json!(finding.stable_id));
250    record.insert("asset_id".into(), json!(finding.asset_id));
251    if let Some(v) = &finding.service_id {
252        record.insert("service_id".into(), json!(v));
253    }
254    record.insert("title".into(), json!(finding.title));
255    record.insert("severity".into(), json!(finding.severity));
256    record.insert("tool".into(), json!(finding.tool));
257    if let Some(v) = &finding.target_url {
258        record.insert("target_url".into(), json!(v));
259    }
260    if let Some(v) = &finding.detail {
261        record.insert("detail".into(), json!(v));
262    }
263    if let Some(v) = &finding.detail_path {
264        record.insert("detail_path".into(), json!(v));
265    }
266    if let Some(v) = &finding.reference {
267        record.insert("reference".into(), json!(v));
268    }
269    record.insert("cves".into(), json!(finding.cves));
270    record.insert("endpoints".into(), json!(finding.endpoints));
271    record.insert("technologies".into(), json!(finding.technologies));
272    record.insert("first_seen".into(), json!(finding.first_seen));
273    record.insert("last_seen".into(), json!(finding.last_seen));
274    record.insert("status".into(), json!(finding.status));
275    if let Some(v) = &finding.exploitability {
276        record.insert("exploitability".into(), json!(v));
277    }
278    record.insert("metadata".into(), as_json_object(&finding.metadata));
279    Value::Object(record)
280}
281
282fn record_to_finding(value: &Value) -> Finding {
283    Finding {
284        stable_id: get_str(value, "stable_id"),
285        asset_id: get_str(value, "asset_id"),
286        service_id: get_opt_str(value, "service_id"),
287        title: get_str(value, "title"),
288        severity: get_str(value, "severity"),
289        tool: get_str(value, "tool"),
290        target_url: get_opt_str(value, "target_url"),
291        detail: get_opt_str(value, "detail"),
292        detail_path: get_opt_str(value, "detail_path"),
293        reference: get_opt_str(value, "reference"),
294        cves: get_vec(value, "cves"),
295        endpoints: get_vec(value, "endpoints"),
296        technologies: get_vec(value, "technologies"),
297        first_seen: get_str(value, "first_seen"),
298        last_seen: get_str(value, "last_seen"),
299        status: get_str(value, "status"),
300        exploitability: get_opt_str(value, "exploitability"),
301        metadata: get_json(value, "metadata"),
302    }
303}
304
305fn observation_to_record(observation: &Observation) -> Value {
306    json!({
307        "stable_id": observation.stable_id,
308        "run_id": observation.run_id,
309        "subject_type": observation.subject_type,
310        "subject_id": observation.subject_id,
311        "kind": observation.kind,
312        "before": as_json_object(&observation.before),
313        "after": as_json_object(&observation.after),
314        "detail": observation.detail,
315        "observed_at": observation.observed_at,
316    })
317}
318
319fn record_to_observation(value: &Value) -> Observation {
320    Observation {
321        stable_id: get_str(value, "stable_id"),
322        run_id: get_str(value, "run_id"),
323        subject_type: get_str(value, "subject_type"),
324        subject_id: get_str(value, "subject_id"),
325        kind: get_str(value, "kind"),
326        before: get_json(value, "before"),
327        after: get_json(value, "after"),
328        detail: get_str(value, "detail"),
329        observed_at: get_str(value, "observed_at"),
330    }
331}
332
333fn scan_run_to_record(run: &ScanRun) -> Value {
334    json!({
335        "stable_id": run.stable_id,
336        "target": run.target,
337        "folder": run.folder,
338        "started_at": run.started_at,
339        "finished_at": run.finished_at,
340        "tools": run.tools,
341        "artifacts": run.artifacts,
342        "content_hash": run.content_hash,
343        "imported_at": run.imported_at,
344        "stats": as_json_object(&run.stats),
345    })
346}
347
348fn record_to_scan_run(value: &Value) -> ScanRun {
349    ScanRun {
350        stable_id: get_str(value, "stable_id"),
351        target: get_str(value, "target"),
352        folder: get_str(value, "folder"),
353        started_at: get_opt_str(value, "started_at"),
354        finished_at: get_opt_str(value, "finished_at"),
355        tools: get_vec(value, "tools"),
356        artifacts: get_vec_value(value, "artifacts"),
357        content_hash: get_str(value, "content_hash"),
358        imported_at: get_str(value, "imported_at"),
359        stats: get_json(value, "stats"),
360    }
361}
362
363fn get_options_map(value: &Value, key: &str) -> std::collections::BTreeMap<String, String> {
364    value
365        .get(key)
366        .and_then(Value::as_object)
367        .map(|map| {
368            map.iter()
369                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
370                .collect()
371        })
372        .unwrap_or_default()
373}
374
375fn exploit_execution_to_record(execution: &ExploitExecution) -> Value {
376    json!({
377        "execution_id": execution.execution_id,
378        "run_id": execution.run_id,
379        "asset_id": execution.asset_id,
380        "vulnerability_id": execution.vulnerability_id,
381        "module_type": execution.module_type,
382        "module": execution.module,
383        "host": execution.host,
384        "payload": execution.payload,
385        "status": execution.status,
386        "start_time": execution.start_time,
387        "finish_time": execution.finish_time,
388        "job_id": execution.job_id,
389        "session_id": execution.session_id,
390        "error": execution.error,
391        "loot": execution.loot,
392        "options": execution.options,
393        "selected": execution.selected,
394        "created_at": execution.created_at,
395    })
396}
397
398fn record_to_exploit_execution(value: &Value) -> ExploitExecution {
399    ExploitExecution {
400        execution_id: get_str(value, "execution_id"),
401        run_id: get_str(value, "run_id"),
402        asset_id: get_str(value, "asset_id"),
403        vulnerability_id: get_str(value, "vulnerability_id"),
404        module_type: get_str(value, "module_type"),
405        module: get_str(value, "module"),
406        host: get_str(value, "host"),
407        payload: get_opt_str(value, "payload"),
408        status: get_str(value, "status"),
409        start_time: get_str(value, "start_time"),
410        finish_time: get_opt_str(value, "finish_time"),
411        job_id: get_opt_str(value, "job_id"),
412        session_id: get_opt_str(value, "session_id"),
413        error: get_opt_str(value, "error"),
414        loot: get_vec(value, "loot"),
415        options: get_options_map(value, "options"),
416        selected: value
417            .get("selected")
418            .and_then(Value::as_bool)
419            .unwrap_or(false),
420        created_at: get_str(value, "created_at"),
421    }
422}
423
424impl PocketBaseStorage {
425    pub fn new(config: &DatabaseConfig) -> Result<Self> {
426        let mut url = config.pocketbase_url.trim().to_string();
427        if url.is_empty() {
428            url = "http://127.0.0.1:8090".to_string();
429        }
430        let url = url.trim_end_matches('/').to_string();
431        Ok(Self {
432            config: DatabaseConfig {
433                pocketbase_url: url,
434                ..config.clone()
435            },
436            token: config.pocketbase_token.clone(),
437            token_validated: false,
438        })
439    }
440
441    fn base_url(&self) -> &str {
442        &self.config.pocketbase_url
443    }
444
445    /// Current auth token, if a login has been performed.
446    pub fn token(&self) -> String {
447        self.token.clone()
448    }
449
450    /// Authenticate with the configured superuser credentials and cache a
451    /// fresh token. Returns the token.
452    pub fn login(&mut self) -> Result<String> {
453        let email = self.config.pocketbase_admin_email.trim();
454        let password = self.config.pocketbase_admin_password.trim();
455        if email.is_empty() || password.is_empty() {
456            return Err(AppError::Message(
457                "PocketBase requires superuser credentials (email + password) to connect. \
458                 Run `ploit-malper db_setup` to configure them."
459                    .to_string(),
460            ));
461        }
462        let url = format!(
463            "{}/api/collections/_superusers/auth-with-password",
464            self.base_url()
465        );
466        let body = json!({ "identity": email, "password": password });
467        let resp = curl_request("POST", &url, None, Some(&body), &[])?;
468        if resp.status != 200 && resp.status != 201 {
469            return Err(AppError::Network(format!(
470                "PocketBase authentication failed (HTTP {}): {}",
471                resp.status, resp.body
472            )));
473        }
474        let token = get_str(&resp.body, "token");
475        if token.is_empty() {
476            return Err(AppError::Network(format!(
477                "PocketBase authentication response did not include a token: {}",
478                resp.body
479            )));
480        }
481        self.token = token.clone();
482        self.config.pocketbase_token = token.clone();
483        self.token_validated = true;
484        Ok(token)
485    }
486
487    /// Ensure we have a usable token, verifying a cached one once and minting
488    /// a fresh token from credentials if it has expired.
489    fn ensure_token(&mut self) -> Result<String> {
490        if self.token_validated && !self.token.is_empty() {
491            return Ok(self.token.clone());
492        }
493        if !self.token.is_empty() {
494            let url = format!("{}/api/collections", self.base_url());
495            let resp = curl_request("GET", &url, Some(&self.token), None, &[("perPage", "1")])?;
496            if resp.status == 200 {
497                self.token_validated = true;
498                return Ok(self.token.clone());
499            }
500            self.token.clear();
501        }
502        self.login()
503    }
504
505    pub fn list_collection_names(&mut self) -> Result<Vec<String>> {
506        let token = self.ensure_token()?;
507        let url = format!("{}/api/collections", self.base_url());
508        let resp = curl_request("GET", &url, Some(&token), None, &[("perPage", "200")])?;
509        if resp.status != 200 {
510            return Err(AppError::Network(format!(
511                "Failed to list PocketBase collections (HTTP {}): {}",
512                resp.status, resp.body
513            )));
514        }
515        let names = resp
516            .body
517            .get("items")
518            .and_then(Value::as_array)
519            .map(|items| {
520                items
521                    .iter()
522                    .filter_map(|c| c.get("name").and_then(Value::as_str).map(str::to_string))
523                    .collect()
524            })
525            .unwrap_or_default();
526        Ok(names)
527    }
528
529    fn find_record(
530        &mut self,
531        collection: &str,
532        stable_id: &str,
533    ) -> Result<Option<(String, Value)>> {
534        let token = self.ensure_token()?;
535        let url = format!("{}/api/collections/{}/records", self.base_url(), collection);
536        let filter = format!("stable_id = \"{}\"", stable_id);
537        let resp = curl_request(
538            "GET",
539            &url,
540            Some(&token),
541            None,
542            &[("filter", &filter), ("perPage", "500")],
543        )?;
544        if resp.status != 200 {
545            return Err(AppError::Network(format!(
546                "Failed to query records in '{}' (HTTP {}): {}",
547                collection, resp.status, resp.body
548            )));
549        }
550        let items = resp
551            .body
552            .get("items")
553            .and_then(Value::as_array)
554            .cloned()
555            .unwrap_or_default();
556        if let Some(first) = items.into_iter().next() {
557            let id = get_str(&first, "id");
558            return Ok(Some((id, first)));
559        }
560        Ok(None)
561    }
562
563    fn create_record(&mut self, collection: &str, record: &Value) -> Result<()> {
564        let token = self.ensure_token()?;
565        let url = format!("{}/api/collections/{}/records", self.base_url(), collection);
566        let resp = curl_request("POST", &url, Some(&token), Some(record), &[])?;
567        if resp.status != 200 && resp.status != 201 {
568            return Err(AppError::Network(format!(
569                "Failed to create record in '{}' (HTTP {}): {}",
570                collection, resp.status, resp.body
571            )));
572        }
573        Ok(())
574    }
575
576    fn update_record(&mut self, collection: &str, record_id: &str, record: &Value) -> Result<()> {
577        let token = self.ensure_token()?;
578        let url = format!(
579            "{}/api/collections/{}/records/{}",
580            self.base_url(),
581            collection,
582            record_id
583        );
584        let resp = curl_request("PATCH", &url, Some(&token), Some(record), &[])?;
585        if resp.status != 200 && resp.status != 204 {
586            return Err(AppError::Network(format!(
587                "Failed to update record in '{}' (HTTP {}): {}",
588                collection, resp.status, resp.body
589            )));
590        }
591        Ok(())
592    }
593
594    fn delete_record_by_id(&mut self, collection: &str, record_id: &str) -> Result<()> {
595        let token = self.ensure_token()?;
596        let url = format!(
597            "{}/api/collections/{}/records/{}",
598            self.base_url(),
599            collection,
600            record_id
601        );
602        let resp = curl_request("DELETE", &url, Some(&token), None, &[])?;
603        if resp.status != 200 && resp.status != 204 {
604            return Err(AppError::Network(format!(
605                "Failed to delete record in '{}' (HTTP {}): {}",
606                collection, resp.status, resp.body
607            )));
608        }
609        Ok(())
610    }
611
612    fn delete_records_by_filter(&mut self, collection: &str, filter: &str) -> Result<()> {
613        let records = self.list_records(collection, Some(filter))?;
614        for record in &records {
615            let record_id = get_str(record, "id");
616            if !record_id.is_empty() {
617                self.delete_record_by_id(collection, &record_id)?;
618            }
619        }
620        Ok(())
621    }
622
623    fn upsert(&mut self, collection: &str, stable_id: &str, record: &Value) -> Result<()> {
624        if let Some((record_id, _)) = self.find_record(collection, stable_id)? {
625            self.update_record(collection, &record_id, record)
626        } else {
627            self.create_record(collection, record)
628        }
629    }
630
631    fn list_records(&mut self, collection: &str, filter: Option<&str>) -> Result<Vec<Value>> {
632        let token = self.ensure_token()?;
633        let url = format!("{}/api/collections/{}/records", self.base_url(), collection);
634        const PER_PAGE: u64 = 500;
635        let mut items: Vec<Value> = Vec::new();
636        let mut page: u64 = 1;
637        loop {
638            let page_str = page.to_string();
639            let per_str = PER_PAGE.to_string();
640            let mut query: Vec<(&str, &str)> =
641                vec![("page", page_str.as_str()), ("perPage", per_str.as_str())];
642            if let Some(f) = filter {
643                query.push(("filter", f));
644            }
645            let resp = curl_request("GET", &url, Some(&token), None, &query)?;
646            if resp.status != 200 {
647                return Err(AppError::Network(format!(
648                    "Failed to list records in '{}' (HTTP {}): {}",
649                    collection, resp.status, resp.body
650                )));
651            }
652            let batch = resp
653                .body
654                .get("items")
655                .and_then(Value::as_array)
656                .cloned()
657                .unwrap_or_default();
658            let count = batch.len();
659            items.extend(batch);
660            if count < PER_PAGE as usize {
661                break;
662            }
663            page += 1;
664        }
665        Ok(items)
666    }
667
668    /// Create a collection directly via the REST API (fallback when pbctl is
669    /// unavailable).
670    fn create_collection_via_api(&mut self, name: &str) -> Result<()> {
671        let token = self.ensure_token()?;
672        let url = format!("{}/api/collections", self.base_url());
673        let defs = schema::field_defs(name);
674        let fields: Vec<Value> = defs.iter().map(|def| self.field_to_json(def)).collect();
675        let body = json!({
676            "name": name,
677            "type": "base",
678            "fields": fields,
679        });
680        let resp = curl_request("POST", &url, Some(&token), Some(&body), &[])?;
681        if resp.status != 200 && resp.status != 201 {
682            return Err(AppError::Network(format!(
683                "Failed to create PocketBase collection '{}' (HTTP {}): {}",
684                name, resp.status, resp.body
685            )));
686        }
687        Ok(())
688    }
689
690    fn field_to_json(&self, def: &FieldDef) -> Value {
691        let mut field = serde_json::Map::new();
692        field.insert("name".to_string(), json!(def.name));
693        field.insert("type".to_string(), json!(def.kind));
694        field.insert("required".to_string(), json!(def.required));
695        field.insert("unique".to_string(), json!(def.unique));
696        field.insert("hidden".to_string(), json!(false));
697        if def.kind == "select" {
698            field.insert(
699                "values".to_string(),
700                json!(["ip", "hostname", "domain", "url"]),
701            );
702            field.insert("maxSelect".to_string(), json!(1));
703        }
704        Value::Object(field)
705    }
706}
707
708impl Storage for PocketBaseStorage {
709    fn kind(&self) -> &str {
710        "pocketbase"
711    }
712
713    fn ping(&mut self) -> Result<bool> {
714        let url = format!("{}/api/health", self.base_url());
715        let resp = curl_request("GET", &url, None, None, &[])?;
716        if resp.status == 200 {
717            return Ok(true);
718        }
719        // Health may be unauthenticated; fall back to an authenticated call.
720        match self.ensure_token() {
721            Ok(token) => {
722                let resp = curl_request("GET", &url, Some(&token), None, &[])?;
723                Ok(resp.status == 200)
724            }
725            Err(e) => Err(e),
726        }
727    }
728
729    fn ensure_schema(&mut self) -> Result<()> {
730        // Pure REST API path — no pbctl required. Create any required
731        // collections that are missing; leave existing ones untouched.
732        let existing = self.list_collection_names()?;
733        for name in collections::ALL {
734            if !existing.iter().any(|n| n == name) {
735                self.create_collection_via_api(name)?;
736            }
737        }
738        Ok(())
739    }
740
741    fn upsert_scan_run(&mut self, run: &ScanRun) -> Result<()> {
742        self.upsert(
743            collections::SCAN_RUNS,
744            &run.stable_id,
745            &scan_run_to_record(run),
746        )
747    }
748
749    fn get_scan_run(&mut self, run_id: &str) -> Result<Option<ScanRun>> {
750        if let Some((_, record)) = self.find_record(collections::SCAN_RUNS, run_id)? {
751            return Ok(Some(record_to_scan_run(&record)));
752        }
753        Ok(None)
754    }
755
756    fn list_scan_runs(&mut self) -> Result<Vec<ScanRun>> {
757        let records = self.list_records(collections::SCAN_RUNS, None)?;
758        Ok(records.iter().map(record_to_scan_run).collect())
759    }
760
761    fn delete_scan_run(&mut self, run_id: &str) -> Result<()> {
762        let obs_filter = format!("run_id = \"{}\"", run_id);
763        self.delete_records_by_filter(collections::OBSERVATIONS, &obs_filter)?;
764        let exec_filter = format!("run_id = \"{}\"", run_id);
765        self.delete_records_by_filter(collections::EXPLOIT_EXECUTIONS, &exec_filter)?;
766        if let Some((record_id, _)) = self.find_record(collections::SCAN_RUNS, run_id)? {
767            self.delete_record_by_id(collections::SCAN_RUNS, &record_id)?;
768        }
769        Ok(())
770    }
771
772    fn upsert_asset(&mut self, asset: &Asset) -> Result<()> {
773        self.upsert(
774            collections::ASSETS,
775            &asset.stable_id,
776            &asset_to_record(asset),
777        )
778    }
779
780    fn get_asset(&mut self, asset_id: &str) -> Result<Option<Asset>> {
781        if let Some((_, record)) = self.find_record(collections::ASSETS, asset_id)? {
782            return Ok(Some(record_to_asset(&record)));
783        }
784        Ok(None)
785    }
786
787    fn list_assets(&mut self) -> Result<Vec<Asset>> {
788        let records = self.list_records(collections::ASSETS, None)?;
789        Ok(records.iter().map(record_to_asset).collect())
790    }
791
792    fn upsert_service(&mut self, service: &Service) -> Result<()> {
793        self.upsert(
794            collections::SERVICES,
795            &service.stable_id,
796            &service_to_record(service),
797        )
798    }
799
800    fn get_service(&mut self, service_id: &str) -> Result<Option<Service>> {
801        if let Some((_, record)) = self.find_record(collections::SERVICES, service_id)? {
802            return Ok(Some(record_to_service(&record)));
803        }
804        Ok(None)
805    }
806
807    fn list_services_for_asset(&mut self, asset_id: &str) -> Result<Vec<Service>> {
808        let filter = format!("asset_id = \"{}\"", asset_id);
809        let records = self.list_records(collections::SERVICES, Some(&filter))?;
810        Ok(records.iter().map(record_to_service).collect())
811    }
812
813    fn list_all_services(&mut self) -> Result<Vec<Service>> {
814        let records = self.list_records(collections::SERVICES, None)?;
815        Ok(records.iter().map(record_to_service).collect())
816    }
817
818    fn upsert_finding(&mut self, finding: &Finding) -> Result<()> {
819        self.upsert(
820            collections::FINDINGS,
821            &finding.stable_id,
822            &finding_to_record(finding),
823        )
824    }
825
826    fn get_finding(&mut self, finding_id: &str) -> Result<Option<Finding>> {
827        if let Some((_, record)) = self.find_record(collections::FINDINGS, finding_id)? {
828            return Ok(Some(record_to_finding(&record)));
829        }
830        Ok(None)
831    }
832
833    fn list_findings_for_asset(&mut self, asset_id: &str) -> Result<Vec<Finding>> {
834        let filter = format!("asset_id = \"{}\"", asset_id);
835        let records = self.list_records(collections::FINDINGS, Some(&filter))?;
836        Ok(records.iter().map(record_to_finding).collect())
837    }
838
839    fn list_all_findings(&mut self) -> Result<Vec<Finding>> {
840        let records = self.list_records(collections::FINDINGS, None)?;
841        Ok(records.iter().map(record_to_finding).collect())
842    }
843
844    fn add_observation(&mut self, observation: &Observation) -> Result<()> {
845        if self.get_observation(&observation.stable_id)?.is_some() {
846            return Ok(());
847        }
848        self.create_record(
849            collections::OBSERVATIONS,
850            &observation_to_record(observation),
851        )
852    }
853
854    fn get_observation(&mut self, obs_id: &str) -> Result<Option<Observation>> {
855        if let Some((_, record)) = self.find_record(collections::OBSERVATIONS, obs_id)? {
856            return Ok(Some(record_to_observation(&record)));
857        }
858        Ok(None)
859    }
860
861    fn list_observations_for(
862        &mut self,
863        subject_type: &str,
864        subject_id: &str,
865    ) -> Result<Vec<Observation>> {
866        let filter = format!(
867            "subject_type = \"{}\" && subject_id = \"{}\"",
868            subject_type, subject_id
869        );
870        let records = self.list_records(collections::OBSERVATIONS, Some(&filter))?;
871        Ok(records.iter().map(record_to_observation).collect())
872    }
873
874    fn list_all_observations(&mut self) -> Result<Vec<Observation>> {
875        let records = self.list_records(collections::OBSERVATIONS, None)?;
876        Ok(records.iter().map(record_to_observation).collect())
877    }
878
879    fn upsert_exploit_execution(&mut self, execution: &ExploitExecution) -> Result<()> {
880        self.upsert(
881            collections::EXPLOIT_EXECUTIONS,
882            &execution.execution_id,
883            &exploit_execution_to_record(execution),
884        )
885    }
886
887    fn get_exploit_execution(&mut self, execution_id: &str) -> Result<Option<ExploitExecution>> {
888        if let Some((_, record)) =
889            self.find_record(collections::EXPLOIT_EXECUTIONS, execution_id)?
890        {
891            return Ok(Some(record_to_exploit_execution(&record)));
892        }
893        Ok(None)
894    }
895
896    fn list_exploit_executions(&mut self) -> Result<Vec<ExploitExecution>> {
897        let records = self.list_records(collections::EXPLOIT_EXECUTIONS, None)?;
898        Ok(records.iter().map(record_to_exploit_execution).collect())
899    }
900
901    fn list_exploit_executions_for_run(&mut self, run_id: &str) -> Result<Vec<ExploitExecution>> {
902        let filter = format!("run_id = \"{}\"", run_id);
903        let records = self.list_records(collections::EXPLOIT_EXECUTIONS, Some(&filter))?;
904        Ok(records.iter().map(record_to_exploit_execution).collect())
905    }
906}
ploit_malper/db/pocketbase.rs

ploit_malper/db/
pocketbase.rs
