1use 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
14impl Indicator {
17 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 pub fn context(mut self, context: impl Into<String>) -> Self {
28 self.context = Some(context.into());
29 self
30 }
31 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 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 pub fn evidence(mut self, ev: impl Into<String>) -> Self {
49 self.evidence = Some(ev.into());
50 self
51 }
52}
53
54impl ArtifactRef {
55 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 pub fn description(mut self, desc: impl Into<String>) -> Self {
65 self.description = Some(desc.into());
66 self
67 }
68}
69
70pub struct ReportBuilder {
74 report: Report,
75 pending_labels: Vec<String>,
76}
77
78impl ReportBuilder {
79 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 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 pub fn summary(mut self, summary: impl Into<String>) -> Self {
110 self.report.summary = Some(summary.into());
111 self
112 }
113
114 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 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 pub fn indicator(mut self, ind: Indicator) -> Self {
145 self.report.indicators.push(ind);
146 self
147 }
148
149 pub fn ttp(mut self, ttp: Ttp) -> Self {
151 self.report.ttps.push(ttp);
152 self
153 }
154
155 pub fn artifact(mut self, artifact: ArtifactRef) -> Self {
157 self.report.artifacts.push(artifact);
158 self
159 }
160
161 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 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 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
197pub 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 pub fn block(mut self, block: Block) -> Self {
218 self.section.blocks.push(block);
219 self
220 }
221
222 pub fn markdown(self, text: impl Into<String>) -> Self {
224 self.block(Block::Markdown { text: text.into() })
225 }
226
227 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 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 pub fn divider(self) -> Self {
245 self.block(Block::Divider)
246 }
247
248 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 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 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 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 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 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 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 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 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 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 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 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}