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

ploit_malper/db/
pocketbase.rs
