1use std::collections::{BTreeMap, BTreeSet};
18
19use rsigma_parser::{
20 CorrelationCondition, CorrelationRule, Detection, DetectionItem, Detections, FilterRule,
21 SigmaCollection, SigmaRule,
22};
23use serde::Serialize;
24
25use crate::pipeline::{Pipeline, apply_pipelines};
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
29#[serde(rename_all = "lowercase")]
30pub enum FieldSource {
31 Detection,
33 Correlation,
35 Filter,
37 Metadata,
39}
40
41impl FieldSource {
42 pub fn as_str(self) -> &'static str {
44 match self {
45 FieldSource::Detection => "detection",
46 FieldSource::Correlation => "correlation",
47 FieldSource::Filter => "filter",
48 FieldSource::Metadata => "metadata",
49 }
50 }
51}
52
53#[derive(Debug, Clone, Default, PartialEq, Eq)]
55pub struct FieldOrigin {
56 pub rule_titles: BTreeSet<String>,
58 pub sources: BTreeSet<FieldSource>,
61}
62
63#[derive(Debug, Clone, Default, PartialEq, Eq)]
70pub struct RuleFieldSet {
71 fields: BTreeMap<String, FieldOrigin>,
72}
73
74impl RuleFieldSet {
75 pub fn collect(
86 collection: &SigmaCollection,
87 pipelines: &[Pipeline],
88 include_filters: bool,
89 ) -> Self {
90 let mut collector = Collector::default();
91
92 if pipelines.is_empty() {
93 for rule in &collection.rules {
94 collector.collect_rule(rule);
95 }
96 } else {
97 for rule in &collection.rules {
98 let mut transformed = rule.clone();
99 if apply_pipelines(pipelines, &mut transformed).is_err() {
100 collector.collect_rule(rule);
101 continue;
102 }
103 collector.collect_rule(&transformed);
104 }
105 }
106
107 for corr in &collection.correlations {
108 collector.collect_correlation(corr);
109 }
110
111 if include_filters {
112 for filter in &collection.filters {
113 collector.collect_filter(filter);
114 }
115 }
116
117 Self {
118 fields: collector.fields,
119 }
120 }
121
122 pub fn contains(&self, field: &str) -> bool {
124 self.fields.contains_key(field)
125 }
126
127 pub fn origin(&self, field: &str) -> Option<&FieldOrigin> {
129 self.fields.get(field)
130 }
131
132 pub fn iter(&self) -> impl Iterator<Item = (&str, &FieldOrigin)> {
134 self.fields.iter().map(|(k, v)| (k.as_str(), v))
135 }
136
137 pub fn names(&self) -> impl Iterator<Item = &str> {
139 self.fields.keys().map(String::as_str)
140 }
141
142 pub fn len(&self) -> usize {
144 self.fields.len()
145 }
146
147 pub fn is_empty(&self) -> bool {
149 self.fields.is_empty()
150 }
151}
152
153#[derive(Default)]
154struct Collector {
155 fields: BTreeMap<String, FieldOrigin>,
156}
157
158impl Collector {
159 fn add(&mut self, field: &str, rule_title: &str, source: FieldSource) {
160 let entry = self.fields.entry(field.to_string()).or_default();
161 entry.rule_titles.insert(rule_title.to_string());
162 entry.sources.insert(source);
163 }
164
165 fn collect_detection_items(
166 &mut self,
167 detection: &Detection,
168 rule_title: &str,
169 source: FieldSource,
170 ) {
171 match detection {
172 Detection::AllOf(items) => {
173 for item in items {
174 self.collect_item(item, rule_title, source);
175 }
176 }
177 Detection::AnyOf(subs) => {
178 for sub in subs {
179 self.collect_detection_items(sub, rule_title, source);
180 }
181 }
182 Detection::ArrayMatch { field, body, .. } => {
183 self.add(field, rule_title, source);
186 self.collect_detection_items(body, rule_title, source);
187 }
188 Detection::And(subs) => {
189 for sub in subs {
190 self.collect_detection_items(sub, rule_title, source);
191 }
192 }
193 Detection::Conditional { named, .. } => {
194 for sub in named.values() {
195 self.collect_detection_items(sub, rule_title, source);
196 }
197 }
198 Detection::Keywords(_) => {}
199 }
200 }
201
202 fn collect_item(&mut self, item: &DetectionItem, rule_title: &str, source: FieldSource) {
203 if let Some(ref name) = item.field.name {
204 self.add(name, rule_title, source);
205 }
206 }
207
208 fn collect_detections(
209 &mut self,
210 detections: &Detections,
211 rule_title: &str,
212 source: FieldSource,
213 ) {
214 for det in detections.named.values() {
215 self.collect_detection_items(det, rule_title, source);
216 }
217 }
218
219 fn collect_rule(&mut self, rule: &SigmaRule) {
220 self.collect_detections(&rule.detection, &rule.title, FieldSource::Detection);
221 for f in &rule.fields {
222 self.add(f, &rule.title, FieldSource::Metadata);
223 }
224 }
225
226 fn collect_correlation(&mut self, corr: &CorrelationRule) {
227 for f in &corr.group_by {
228 self.add(f, &corr.title, FieldSource::Correlation);
229 }
230 if let CorrelationCondition::Threshold {
231 field: Some(ref fields),
232 ..
233 } = corr.condition
234 {
235 for f in fields {
236 self.add(f, &corr.title, FieldSource::Correlation);
237 }
238 }
239 for alias in &corr.aliases {
240 for mapped_field in alias.mapping.values() {
241 self.add(mapped_field, &corr.title, FieldSource::Correlation);
242 }
243 }
244 for f in &corr.fields {
245 self.add(f, &corr.title, FieldSource::Metadata);
246 }
247 }
248
249 fn collect_filter(&mut self, filter: &FilterRule) {
250 self.collect_detections(&filter.detection, &filter.title, FieldSource::Filter);
251 for f in &filter.fields {
252 self.add(f, &filter.title, FieldSource::Metadata);
253 }
254 }
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260 use rsigma_parser::parse_sigma_yaml;
261
262 fn build(yaml: &str) -> SigmaCollection {
263 parse_sigma_yaml(yaml).expect("parse")
264 }
265
266 #[test]
267 fn collects_detection_fields() {
268 let collection = build(
269 r#"
270title: Test
271status: test
272logsource:
273 category: test
274detection:
275 selection:
276 CommandLine|contains: whoami
277 EventID: 1
278 condition: selection
279"#,
280 );
281 let set = RuleFieldSet::collect(&collection, &[], true);
282 assert!(set.contains("CommandLine"));
283 assert!(set.contains("EventID"));
284 assert!(
285 set.origin("CommandLine")
286 .unwrap()
287 .sources
288 .contains(&FieldSource::Detection)
289 );
290 }
291
292 #[test]
293 fn collects_correlation_group_by() {
294 let collection = build(
295 r#"
296title: Login
297id: login-rule
298logsource:
299 category: auth
300detection:
301 selection:
302 EventType: login
303 condition: selection
304---
305title: Many Logins
306correlation:
307 type: event_count
308 rules:
309 - login-rule
310 group-by:
311 - User
312 timespan: 60s
313 condition:
314 gte: 3
315"#,
316 );
317 let set = RuleFieldSet::collect(&collection, &[], true);
318 assert!(set.contains("EventType"));
319 assert!(set.contains("User"));
320 let user_origin = set.origin("User").unwrap();
321 assert!(user_origin.sources.contains(&FieldSource::Correlation));
322 }
323
324 #[test]
325 fn include_filters_toggle() {
326 let collection = build(
327 r#"
328title: Detection
329status: test
330logsource:
331 category: test
332detection:
333 selection:
334 DetField: x
335 condition: selection
336---
337title: Filter
338filter:
339 rules:
340 - non-existent
341 selection:
342 FilterField: y
343 condition: selection
344"#,
345 );
346 let with_filters = RuleFieldSet::collect(&collection, &[], true);
347 let without_filters = RuleFieldSet::collect(&collection, &[], false);
348 assert!(with_filters.contains("FilterField"));
349 assert!(!without_filters.contains("FilterField"));
350 assert!(with_filters.contains("DetField"));
351 assert!(without_filters.contains("DetField"));
352 }
353
354 #[test]
355 fn empty_collection_is_empty_set() {
356 let collection = SigmaCollection::default();
357 let set = RuleFieldSet::collect(&collection, &[], true);
358 assert!(set.is_empty());
359 assert_eq!(set.len(), 0);
360 }
361}