Skip to main content

malbox_plugin_sdk/report/
builder.rs

1//! Fluent builder API for constructing [`Report`]s.
2//!
3//! Use [`ReportBuilder`] to assemble a report from its parts and
4//! [`SectionBuilder`] (via [`ReportBuilder::section`]) to add presentation
5//! blocks. See the parent [`report`](super) module for how the envelope
6//! is structured.
7
8use super::{
9    ArtifactRef, Block, CalloutLevel, Classification, Column, Confidence, GraphEdge, GraphNode,
10    Indicator, KvPair, PluginInfo, Report, SCHEMA_VERSION, Section, TimelineEvent, TreeNode, Ttp,
11    Verdict,
12};
13
14// ---------- Small constructors on the leaf types ----------
15
16impl Indicator {
17    /// Create an indicator with a type and value (e.g. `"sha256"`, `"abcd1234..."`).
18    pub fn new(kind: impl Into<String>, value: impl Into<String>) -> Self {
19        Self {
20            kind: kind.into(),
21            value: value.into(),
22            context: None,
23            first_seen: None,
24        }
25    }
26    /// Attach context describing where this IOC was observed.
27    pub fn context(mut self, context: impl Into<String>) -> Self {
28        self.context = Some(context.into());
29        self
30    }
31    /// Set the timestamp when this IOC was first observed.
32    pub fn first_seen(mut self, ts: impl Into<String>) -> Self {
33        self.first_seen = Some(ts.into());
34        self
35    }
36}
37
38impl Ttp {
39    /// Create a TTP with a MITRE ATT&CK ID and name.
40    pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
41        Self {
42            id: id.into(),
43            name: name.into(),
44            evidence: None,
45        }
46    }
47    /// Attach free-text evidence supporting this observation.
48    pub fn evidence(mut self, ev: impl Into<String>) -> Self {
49        self.evidence = Some(ev.into());
50        self
51    }
52}
53
54impl ArtifactRef {
55    /// Create a reference to a sibling result by name and artifact type.
56    pub fn new(result_name: impl Into<String>, kind: impl Into<String>) -> Self {
57        Self {
58            result_name: result_name.into(),
59            kind: kind.into(),
60            description: None,
61        }
62    }
63    /// Add a human-readable description of the artifact.
64    pub fn description(mut self, desc: impl Into<String>) -> Self {
65        self.description = Some(desc.into());
66        self
67    }
68}
69
70// ---------- ReportBuilder ----------
71
72/// Fluent builder for a [`Report`]. Finalise with [`ReportBuilder::build`].
73pub struct ReportBuilder {
74    report: Report,
75    pending_labels: Vec<String>,
76}
77
78impl ReportBuilder {
79    /// Start a new report for a plugin. `id` and `version` typically come
80    /// from `env!("CARGO_PKG_NAME")` / `env!("CARGO_PKG_VERSION")`.
81    pub fn new(plugin_id: impl Into<String>, plugin_version: impl Into<String>) -> Self {
82        Self {
83            report: Report {
84                schema_version: SCHEMA_VERSION,
85                plugin: PluginInfo {
86                    id: plugin_id.into(),
87                    version: plugin_version.into(),
88                    display_name: None,
89                },
90                verdict: None,
91                indicators: Vec::new(),
92                ttps: Vec::new(),
93                artifacts: Vec::new(),
94                summary: None,
95                sections: Vec::new(),
96                raw: None,
97            },
98            pending_labels: Vec::new(),
99        }
100    }
101
102    /// Set a human-friendly plugin name for the frontend (e.g. `"YARA Scanner"`).
103    pub fn display_name(mut self, name: impl Into<String>) -> Self {
104        self.report.plugin.display_name = Some(name.into());
105        self
106    }
107
108    /// Set a short summary shown at the top of the report.
109    pub fn summary(mut self, summary: impl Into<String>) -> Self {
110        self.report.summary = Some(summary.into());
111        self
112    }
113
114    /// Set the verdict. `score` is optional (pass `None` when the plugin
115    /// has no numeric score). Can be called before or after [`labels`](Self::labels).
116    pub fn verdict(
117        mut self,
118        classification: Classification,
119        score: Option<u8>,
120        confidence: Option<Confidence>,
121    ) -> Self {
122        self.report.verdict = Some(Verdict {
123            classification,
124            score,
125            confidence,
126            labels: Vec::new(),
127        });
128        self
129    }
130
131    /// Append labels to the verdict. Can be called before or after
132    /// [`verdict`](Self::verdict) - labels are merged at [`build`](Self::build) time.
133    pub fn labels<I, S>(mut self, labels: I) -> Self
134    where
135        I: IntoIterator<Item = S>,
136        S: Into<String>,
137    {
138        self.pending_labels
139            .extend(labels.into_iter().map(Into::into));
140        self
141    }
142
143    /// Add an indicator of compromise to the report's semantic layer.
144    pub fn indicator(mut self, ind: Indicator) -> Self {
145        self.report.indicators.push(ind);
146        self
147    }
148
149    /// Add a MITRE ATT&CK technique observation.
150    pub fn ttp(mut self, ttp: Ttp) -> Self {
151        self.report.ttps.push(ttp);
152        self
153    }
154
155    /// Register a sibling [`PluginResult`](crate::result::PluginResult) as an artifact.
156    pub fn artifact(mut self, artifact: ArtifactRef) -> Self {
157        self.report.artifacts.push(artifact);
158        self
159    }
160
161    /// Append a section. The closure receives a [`SectionBuilder`] and
162    /// returns it - the typical shape is `|s| s.heading(2, "...").table(...)`.
163    pub fn section<F>(mut self, id: impl Into<String>, title: impl Into<String>, build: F) -> Self
164    where
165        F: FnOnce(SectionBuilder) -> SectionBuilder,
166    {
167        let sb = SectionBuilder::new(id, title);
168        self.report.sections.push(build(sb).build());
169        self
170    }
171
172    /// Attach the plugin's native JSON as an escape hatch. Rarely needed -
173    /// prefer typed blocks.
174    pub fn raw(mut self, value: impl serde::Serialize) -> Self {
175        match serde_json::to_value(value) {
176            Ok(v) => self.report.raw = Some(v),
177            Err(e) => tracing::warn!("ReportBuilder::raw: serialization failed: {e}"),
178        }
179        self
180    }
181
182    /// Finalize and return the [`Report`]. Merges any pending labels into the verdict.
183    pub fn build(mut self) -> Report {
184        if !self.pending_labels.is_empty() {
185            let verdict = self.report.verdict.get_or_insert_with(|| Verdict {
186                classification: Classification::Unknown,
187                score: None,
188                confidence: None,
189                labels: Vec::new(),
190            });
191            verdict.labels.extend(self.pending_labels);
192        }
193        self.report
194    }
195}
196
197// ---------- SectionBuilder ----------
198
199/// Builder for a single [`Section`]. One method per [`Block`] variant, plus
200/// a `block(...)` escape hatch for forward compatibility.
201pub struct SectionBuilder {
202    section: Section,
203}
204
205impl SectionBuilder {
206    pub(crate) fn new(id: impl Into<String>, title: impl Into<String>) -> Self {
207        Self {
208            section: Section {
209                id: id.into(),
210                title: title.into(),
211                blocks: Vec::new(),
212            },
213        }
214    }
215
216    /// Escape hatch - push any [`Block`] directly.
217    pub fn block(mut self, block: Block) -> Self {
218        self.section.blocks.push(block);
219        self
220    }
221
222    /// Add a markdown text block.
223    pub fn markdown(self, text: impl Into<String>) -> Self {
224        self.block(Block::Markdown { text: text.into() })
225    }
226
227    /// Add a highlighted callout box with a severity level.
228    pub fn callout(self, level: CalloutLevel, text: impl Into<String>) -> Self {
229        self.block(Block::Callout {
230            level,
231            text: text.into(),
232        })
233    }
234
235    /// Add a heading (level 1-6, maps to HTML heading levels).
236    pub fn heading(self, level: u8, text: impl Into<String>) -> Self {
237        self.block(Block::Heading {
238            level,
239            text: text.into(),
240        })
241    }
242
243    /// Add a horizontal divider.
244    pub fn divider(self) -> Self {
245        self.block(Block::Divider)
246    }
247
248    /// Add a key-value list.
249    pub fn kv(self, pairs: impl IntoIterator<Item = KvPair>) -> Self {
250        self.block(Block::Kv {
251            pairs: pairs.into_iter().collect(),
252        })
253    }
254
255    /// Add a data table. Sortable by default, not searchable.
256    pub fn table(
257        self,
258        columns: impl IntoIterator<Item = Column>,
259        rows: impl IntoIterator<Item = serde_json::Value>,
260    ) -> Self {
261        self.block(Block::Table {
262            columns: columns.into_iter().collect(),
263            rows: rows.into_iter().collect(),
264            sortable: true,
265            searchable: false,
266        })
267    }
268
269    /// Add a syntax-highlighted code block.
270    pub fn code(self, language: impl Into<String>, text: impl Into<String>) -> Self {
271        self.block(Block::Code {
272            language: language.into(),
273            text: text.into(),
274        })
275    }
276
277    /// Add an interactive JSON tree viewer (collapsed by default).
278    pub fn json(self, data: impl serde::Serialize) -> Self {
279        let v = match serde_json::to_value(data) {
280            Ok(v) => v,
281            Err(e) => {
282                tracing::warn!("SectionBuilder::json: serialization failed: {e}");
283                serde_json::Value::Null
284            }
285        };
286        self.block(Block::Json {
287            data: v,
288            collapsed: true,
289        })
290    }
291
292    /// Add a hex dump. `bytes_b64` is the data as base64, `offset` is the starting address.
293    pub fn hex(self, bytes_b64: impl Into<String>, offset: u64) -> Self {
294        self.block(Block::Hex {
295            bytes_b64: bytes_b64.into(),
296            offset,
297        })
298    }
299
300    /// Add an inline image resolved from a sibling artifact result.
301    pub fn image(self, artifact: impl Into<String>, caption: Option<String>) -> Self {
302        self.block(Block::Image {
303            artifact: artifact.into(),
304            caption,
305        })
306    }
307
308    /// Add a download link resolved from a sibling artifact result.
309    pub fn download(self, artifact: impl Into<String>, label: impl Into<String>) -> Self {
310        self.block(Block::Download {
311            artifact: artifact.into(),
312            label: label.into(),
313        })
314    }
315
316    /// Add a formatted IOC list.
317    pub fn iocs(self, items: impl IntoIterator<Item = Indicator>) -> Self {
318        self.block(Block::Iocs {
319            items: items.into_iter().collect(),
320        })
321    }
322
323    /// Add a formatted MITRE ATT&CK technique list.
324    pub fn ttps(self, items: impl IntoIterator<Item = Ttp>) -> Self {
325        self.block(Block::Ttps {
326            items: items.into_iter().collect(),
327        })
328    }
329
330    /// Add a collapsible tree (e.g. process tree, file hierarchy).
331    pub fn tree(self, nodes: impl IntoIterator<Item = TreeNode>) -> Self {
332        self.block(Block::Tree {
333            nodes: nodes.into_iter().collect(),
334        })
335    }
336
337    /// Add a chronological event timeline.
338    pub fn timeline(self, events: impl IntoIterator<Item = TimelineEvent>) -> Self {
339        self.block(Block::Timeline {
340            events: events.into_iter().collect(),
341        })
342    }
343
344    /// Add a node-and-edge graph (e.g. network map, call graph).
345    pub fn graph(
346        self,
347        nodes: impl IntoIterator<Item = GraphNode>,
348        edges: impl IntoIterator<Item = GraphEdge>,
349    ) -> Self {
350        self.block(Block::Graph {
351            nodes: nodes.into_iter().collect(),
352            edges: edges.into_iter().collect(),
353        })
354    }
355
356    pub(crate) fn build(self) -> Section {
357        self.section
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use serde_json::json;
365
366    #[test]
367    fn builder_produces_minimal_report() {
368        let r = ReportBuilder::new("yara", "1.0.0").build();
369        assert_eq!(r.schema_version, SCHEMA_VERSION);
370        assert_eq!(r.plugin.id, "yara");
371        assert!(r.verdict.is_none());
372        assert!(r.sections.is_empty());
373    }
374
375    #[test]
376    fn builder_sets_verdict_and_labels() {
377        let r = ReportBuilder::new("yara", "1.0.0")
378            .verdict(Classification::Malicious, Some(87), Some(Confidence::High))
379            .labels(["trojan", "stealer"])
380            .build();
381        let v = r.verdict.unwrap();
382        assert_eq!(v.classification, Classification::Malicious);
383        assert_eq!(v.score, Some(87));
384        assert_eq!(v.labels, vec!["trojan", "stealer"]);
385    }
386
387    #[test]
388    fn labels_before_verdict_still_works() {
389        let r = ReportBuilder::new("p", "0")
390            .labels(["tagA"])
391            .verdict(Classification::Suspicious, None, None)
392            .build();
393        let v = r.verdict.unwrap();
394        assert_eq!(v.classification, Classification::Suspicious);
395        assert_eq!(v.labels, vec!["tagA"]);
396    }
397
398    #[test]
399    fn section_builder_adds_blocks_in_order() {
400        let r = ReportBuilder::new("yara", "1.0.0")
401            .section("overview", "Overview", |s| {
402                s.heading(2, "Rules")
403                    .markdown("1 match")
404                    .code("yara", "rule r { condition: true }")
405                    .divider()
406            })
407            .build();
408
409        let sec = &r.sections[0];
410        assert_eq!(sec.id, "overview");
411        assert_eq!(sec.blocks.len(), 4);
412        assert!(matches!(sec.blocks[0], Block::Heading { .. }));
413        assert!(matches!(sec.blocks[3], Block::Divider));
414    }
415
416    #[test]
417    fn indicator_and_ttp_helpers_work() {
418        let r = ReportBuilder::new("p", "0")
419            .indicator(Indicator::new("sha256", "abc").context("sample"))
420            .indicator(Indicator::new("ipv4", "1.2.3.4"))
421            .ttp(Ttp::new("T1055", "Process Injection").evidence("memdump"))
422            .artifact(ArtifactRef::new("cap.pcap", "pcap").description("Full capture"))
423            .build();
424
425        assert_eq!(r.indicators.len(), 2);
426        assert_eq!(r.indicators[0].context.as_deref(), Some("sample"));
427        assert_eq!(r.ttps[0].evidence.as_deref(), Some("memdump"));
428        assert_eq!(r.artifacts[0].kind, "pcap");
429    }
430
431    #[test]
432    fn raw_escape_hatch_stores_json() {
433        #[derive(serde::Serialize)]
434        struct Native {
435            hit: u32,
436        }
437        let r = ReportBuilder::new("p", "0").raw(Native { hit: 3 }).build();
438        assert_eq!(r.raw.unwrap(), json!({"hit": 3}));
439    }
440}