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

ploit_malper/db/
pocketbase.rs
