Skip to main content

ploit_malper/db/
schema.rs

1use std::process::Command;
2
3use crate::config::{db_schema_file, DatabaseConfig};
4use crate::error::{AppError, Result};
5
6/// Canonical PBSL schema shipped with PloitMalper. Written to the config
7/// directory so `pbctl` can be pointed at it.
8pub const SCHEMA_PBSL: &str = include_str!("../../schema/schema.pbsl");
9
10/// Collection names and their field definitions used for the direct
11/// PocketBase REST fallback when `pbctl` is not available.
12pub const REQUIRED_COLLECTIONS: [&str; 7] = [
13    "Assets",
14    "Services",
15    "Findings",
16    "Observations",
17    "Relationships",
18    "ScanRuns",
19    "Reports",
20];
21
22pub struct FieldDef {
23    pub name: &'static str,
24    pub kind: &'static str,
25    pub required: bool,
26    pub unique: bool,
27}
28
29pub fn field_defs(collection: &str) -> &'static [FieldDef] {
30    match collection {
31        "Assets" => &[
32            FieldDef {
33                name: "stable_id",
34                kind: "text",
35                required: true,
36                unique: true,
37            },
38            FieldDef {
39                name: "asset_type",
40                kind: "select",
41                required: false,
42                unique: false,
43            },
44            FieldDef {
45                name: "name",
46                kind: "text",
47                required: true,
48                unique: false,
49            },
50            FieldDef {
51                name: "ip",
52                kind: "text",
53                required: false,
54                unique: false,
55            },
56            FieldDef {
57                name: "fqdn",
58                kind: "text",
59                required: false,
60                unique: false,
61            },
62            FieldDef {
63                name: "reverse_dns",
64                kind: "text",
65                required: false,
66                unique: false,
67            },
68            FieldDef {
69                name: "first_seen",
70                kind: "text",
71                required: false,
72                unique: false,
73            },
74            FieldDef {
75                name: "last_seen",
76                kind: "text",
77                required: false,
78                unique: false,
79            },
80            FieldDef {
81                name: "status",
82                kind: "text",
83                required: false,
84                unique: false,
85            },
86            FieldDef {
87                name: "metadata",
88                kind: "json",
89                required: false,
90                unique: false,
91            },
92        ],
93        "Services" => &[
94            FieldDef {
95                name: "stable_id",
96                kind: "text",
97                required: true,
98                unique: true,
99            },
100            FieldDef {
101                name: "asset_id",
102                kind: "text",
103                required: true,
104                unique: false,
105            },
106            FieldDef {
107                name: "port",
108                kind: "number",
109                required: false,
110                unique: false,
111            },
112            FieldDef {
113                name: "protocol",
114                kind: "text",
115                required: false,
116                unique: false,
117            },
118            FieldDef {
119                name: "service_name",
120                kind: "text",
121                required: false,
122                unique: false,
123            },
124            FieldDef {
125                name: "product",
126                kind: "text",
127                required: false,
128                unique: false,
129            },
130            FieldDef {
131                name: "version",
132                kind: "text",
133                required: false,
134                unique: false,
135            },
136            FieldDef {
137                name: "version_str",
138                kind: "text",
139                required: false,
140                unique: false,
141            },
142            FieldDef {
143                name: "banner",
144                kind: "text",
145                required: false,
146                unique: false,
147            },
148            FieldDef {
149                name: "technologies",
150                kind: "json",
151                required: false,
152                unique: false,
153            },
154            FieldDef {
155                name: "cpes",
156                kind: "json",
157                required: false,
158                unique: false,
159            },
160            FieldDef {
161                name: "first_seen",
162                kind: "text",
163                required: false,
164                unique: false,
165            },
166            FieldDef {
167                name: "last_seen",
168                kind: "text",
169                required: false,
170                unique: false,
171            },
172            FieldDef {
173                name: "status",
174                kind: "text",
175                required: false,
176                unique: false,
177            },
178            FieldDef {
179                name: "metadata",
180                kind: "json",
181                required: false,
182                unique: false,
183            },
184        ],
185        "Findings" => &[
186            FieldDef {
187                name: "stable_id",
188                kind: "text",
189                required: true,
190                unique: true,
191            },
192            FieldDef {
193                name: "asset_id",
194                kind: "text",
195                required: true,
196                unique: false,
197            },
198            FieldDef {
199                name: "service_id",
200                kind: "text",
201                required: false,
202                unique: false,
203            },
204            FieldDef {
205                name: "title",
206                kind: "text",
207                required: true,
208                unique: false,
209            },
210            FieldDef {
211                name: "severity",
212                kind: "text",
213                required: false,
214                unique: false,
215            },
216            FieldDef {
217                name: "tool",
218                kind: "text",
219                required: false,
220                unique: false,
221            },
222            FieldDef {
223                name: "target_url",
224                kind: "text",
225                required: false,
226                unique: false,
227            },
228            FieldDef {
229                name: "detail",
230                kind: "text",
231                required: false,
232                unique: false,
233            },
234            FieldDef {
235                name: "reference",
236                kind: "text",
237                required: false,
238                unique: false,
239            },
240            FieldDef {
241                name: "cves",
242                kind: "json",
243                required: false,
244                unique: false,
245            },
246            FieldDef {
247                name: "endpoints",
248                kind: "json",
249                required: false,
250                unique: false,
251            },
252            FieldDef {
253                name: "technologies",
254                kind: "json",
255                required: false,
256                unique: false,
257            },
258            FieldDef {
259                name: "first_seen",
260                kind: "text",
261                required: false,
262                unique: false,
263            },
264            FieldDef {
265                name: "last_seen",
266                kind: "text",
267                required: false,
268                unique: false,
269            },
270            FieldDef {
271                name: "status",
272                kind: "text",
273                required: false,
274                unique: false,
275            },
276            FieldDef {
277                name: "exploitability",
278                kind: "text",
279                required: false,
280                unique: false,
281            },
282            FieldDef {
283                name: "metadata",
284                kind: "json",
285                required: false,
286                unique: false,
287            },
288        ],
289        "Observations" => &[
290            FieldDef {
291                name: "stable_id",
292                kind: "text",
293                required: true,
294                unique: true,
295            },
296            FieldDef {
297                name: "run_id",
298                kind: "text",
299                required: true,
300                unique: false,
301            },
302            FieldDef {
303                name: "subject_type",
304                kind: "text",
305                required: true,
306                unique: false,
307            },
308            FieldDef {
309                name: "subject_id",
310                kind: "text",
311                required: true,
312                unique: false,
313            },
314            FieldDef {
315                name: "kind",
316                kind: "text",
317                required: true,
318                unique: false,
319            },
320            FieldDef {
321                name: "before",
322                kind: "json",
323                required: false,
324                unique: false,
325            },
326            FieldDef {
327                name: "after",
328                kind: "json",
329                required: false,
330                unique: false,
331            },
332            FieldDef {
333                name: "detail",
334                kind: "text",
335                required: false,
336                unique: false,
337            },
338            FieldDef {
339                name: "observed_at",
340                kind: "text",
341                required: false,
342                unique: false,
343            },
344        ],
345        "Relationships" => &[
346            FieldDef {
347                name: "stable_id",
348                kind: "text",
349                required: true,
350                unique: true,
351            },
352            FieldDef {
353                name: "run_id",
354                kind: "text",
355                required: false,
356                unique: false,
357            },
358            FieldDef {
359                name: "subject_type",
360                kind: "text",
361                required: true,
362                unique: false,
363            },
364            FieldDef {
365                name: "subject_id",
366                kind: "text",
367                required: true,
368                unique: false,
369            },
370            FieldDef {
371                name: "object_type",
372                kind: "text",
373                required: true,
374                unique: false,
375            },
376            FieldDef {
377                name: "object_id",
378                kind: "text",
379                required: true,
380                unique: false,
381            },
382            FieldDef {
383                name: "kind",
384                kind: "text",
385                required: true,
386                unique: false,
387            },
388            FieldDef {
389                name: "detail",
390                kind: "text",
391                required: false,
392                unique: false,
393            },
394        ],
395        "ScanRuns" => &[
396            FieldDef {
397                name: "stable_id",
398                kind: "text",
399                required: true,
400                unique: true,
401            },
402            FieldDef {
403                name: "target",
404                kind: "text",
405                required: true,
406                unique: false,
407            },
408            FieldDef {
409                name: "folder",
410                kind: "text",
411                required: false,
412                unique: false,
413            },
414            FieldDef {
415                name: "started_at",
416                kind: "text",
417                required: false,
418                unique: false,
419            },
420            FieldDef {
421                name: "finished_at",
422                kind: "text",
423                required: false,
424                unique: false,
425            },
426            FieldDef {
427                name: "tools",
428                kind: "json",
429                required: false,
430                unique: false,
431            },
432            FieldDef {
433                name: "artifacts",
434                kind: "json",
435                required: false,
436                unique: false,
437            },
438            FieldDef {
439                name: "content_hash",
440                kind: "text",
441                required: false,
442                unique: false,
443            },
444            FieldDef {
445                name: "imported_at",
446                kind: "text",
447                required: false,
448                unique: false,
449            },
450            FieldDef {
451                name: "stats",
452                kind: "json",
453                required: false,
454                unique: false,
455            },
456        ],
457        "Reports" => &[
458            FieldDef {
459                name: "stable_id",
460                kind: "text",
461                required: true,
462                unique: true,
463            },
464            FieldDef {
465                name: "run_id",
466                kind: "text",
467                required: true,
468                unique: false,
469            },
470            FieldDef {
471                name: "tool",
472                kind: "text",
473                required: true,
474                unique: false,
475            },
476            FieldDef {
477                name: "format",
478                kind: "text",
479                required: true,
480                unique: false,
481            },
482            FieldDef {
483                name: "title",
484                kind: "text",
485                required: false,
486                unique: false,
487            },
488            FieldDef {
489                name: "source_path",
490                kind: "text",
491                required: false,
492                unique: false,
493            },
494            FieldDef {
495                name: "content",
496                kind: "text",
497                required: false,
498                unique: false,
499            },
500            FieldDef {
501                name: "content_hash",
502                kind: "text",
503                required: false,
504                unique: false,
505            },
506            FieldDef {
507                name: "imported_at",
508                kind: "text",
509                required: false,
510                unique: false,
511            },
512        ],
513        _ => &[],
514    }
515}
516
517/// Write the canonical PBSL schema into the config directory and return its
518/// path.
519pub fn write_schema_file() -> Result<std::path::PathBuf> {
520    let path = db_schema_file();
521    if let Some(parent) = path.parent() {
522        std::fs::create_dir_all(parent)?;
523    }
524    std::fs::write(&path, SCHEMA_PBSL)?;
525    Ok(path)
526}
527
528/// Whether the `pbctl` binary is available on PATH.
529pub fn pbctl_available() -> bool {
530    Command::new("pbctl")
531        .arg("--version")
532        .output()
533        .map(|o| o.status.success())
534        .unwrap_or(false)
535}
536
537/// Combined output of an external command.
538#[derive(Debug, Clone)]
539pub struct CommandOutput {
540    pub stdout: String,
541    pub stderr: String,
542    pub exit_code: Option<i32>,
543}
544
545impl CommandOutput {
546    pub fn combined(&self) -> String {
547        let mut out = self.stdout.clone();
548        if !self.stderr.trim().is_empty() {
549            if !out.is_empty() {
550                out.push('\n');
551            }
552            out.push_str(&self.stderr);
553        }
554        out
555    }
556
557    pub fn success(&self) -> bool {
558        self.exit_code == Some(0)
559    }
560}
561
562/// Build a `pbctl` command with environment overrides so it talks to the
563/// configured instance without needing an interactive login.
564fn pbctl_command(config: &DatabaseConfig) -> Command {
565    let mut command = Command::new("pbctl");
566    command
567        .env("PBSL_BASE_URL", &config.pocketbase_url)
568        .env("PBSL_ADMIN_EMAIL", &config.pocketbase_admin_email)
569        .env("PBSL_ADMIN_PASSWORD", &config.pocketbase_admin_password);
570    if !config.pocketbase_token.is_empty() {
571        command.env("PBSL_TOKEN", &config.pocketbase_token);
572    }
573    command
574}
575
576fn run(command: &mut Command, verbose: bool) -> Result<CommandOutput> {
577    let output = command
578        .output()
579        .map_err(|e| AppError::Network(format!("failed to run pbctl: {e}")))?;
580    let stdout = String::from_utf8_lossy(&output.stdout).to_string();
581    let stderr = String::from_utf8_lossy(&output.stderr).to_string();
582    if verbose {
583        if !stdout.is_empty() {
584            println!("{}", stdout);
585        }
586        if !stderr.is_empty() {
587            eprintln!("{}", stderr);
588        }
589    }
590    Ok(CommandOutput {
591        stdout,
592        stderr,
593        exit_code: output.status.code(),
594    })
595}
596
597/// Run `pbctl plan` against the shipped PBSL schema. Never mutates the
598/// remote instance. Returns the full plan output for the caller to display.
599pub fn pbctl_plan(config: &DatabaseConfig, verbose: bool) -> Result<CommandOutput> {
600    if !pbctl_available() {
601        return Err(AppError::Message(
602            "pbctl is not installed. Install it with `cargo install pbctl`.".to_string(),
603        ));
604    }
605    let schema_path = write_schema_file()?;
606    let mut command = pbctl_command(config);
607    command.arg("plan").arg("-f").arg(&schema_path);
608    run(&mut command, verbose)
609}
610
611/// Run `pbctl apply` against the shipped PBSL schema. Only called after the
612/// user has reviewed the plan output from `pbctl_plan`.
613pub fn pbctl_apply(config: &DatabaseConfig, verbose: bool) -> Result<CommandOutput> {
614    if !pbctl_available() {
615        return Err(AppError::Message(
616            "pbctl is not installed. Install it with `cargo install pbctl`.".to_string(),
617        ));
618    }
619    let schema_path = write_schema_file()?;
620    let mut command = pbctl_command(config);
621    command.arg("apply").arg("-f").arg(&schema_path);
622    run(&mut command, verbose)
623}
624
625/// Run `pbctl validate` to sanity-check the shipped schema file locally.
626pub fn pbctl_validate(verbose: bool) -> Result<CommandOutput> {
627    if !pbctl_available() {
628        return Err(AppError::Message(
629            "pbctl is not installed. Install it with `cargo install pbctl`.".to_string(),
630        ));
631    }
632    let schema_path = write_schema_file()?;
633    let mut command = Command::new("pbctl");
634    command.arg("validate").arg("-f").arg(&schema_path);
635    run(&mut command, verbose)
636}