use std::collections::{BTreeMap, BTreeSet};
use rsigma_parser::{
CorrelationCondition, CorrelationRule, Detection, DetectionItem, Detections, FilterRule,
SigmaCollection, SigmaRule,
};
use serde::Serialize;
use crate::pipeline::{Pipeline, apply_pipelines};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum FieldSource {
Detection,
Correlation,
Filter,
Metadata,
}
impl FieldSource {
pub fn as_str(self) -> &'static str {
match self {
FieldSource::Detection => "detection",
FieldSource::Correlation => "correlation",
FieldSource::Filter => "filter",
FieldSource::Metadata => "metadata",
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct FieldOrigin {
pub rule_titles: BTreeSet<String>,
pub sources: BTreeSet<FieldSource>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuleFieldSet {
fields: BTreeMap<String, FieldOrigin>,
}
impl RuleFieldSet {
pub fn collect(
collection: &SigmaCollection,
pipelines: &[Pipeline],
include_filters: bool,
) -> Self {
let mut collector = Collector::default();
if pipelines.is_empty() {
for rule in &collection.rules {
collector.collect_rule(rule);
}
} else {
for rule in &collection.rules {
let mut transformed = rule.clone();
if apply_pipelines(pipelines, &mut transformed).is_err() {
collector.collect_rule(rule);
continue;
}
collector.collect_rule(&transformed);
}
}
for corr in &collection.correlations {
collector.collect_correlation(corr);
}
if include_filters {
for filter in &collection.filters {
collector.collect_filter(filter);
}
}
Self {
fields: collector.fields,
}
}
pub fn contains(&self, field: &str) -> bool {
self.fields.contains_key(field)
}
pub fn origin(&self, field: &str) -> Option<&FieldOrigin> {
self.fields.get(field)
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &FieldOrigin)> {
self.fields.iter().map(|(k, v)| (k.as_str(), v))
}
pub fn names(&self) -> impl Iterator<Item = &str> {
self.fields.keys().map(String::as_str)
}
pub fn len(&self) -> usize {
self.fields.len()
}
pub fn is_empty(&self) -> bool {
self.fields.is_empty()
}
}
#[derive(Default)]
struct Collector {
fields: BTreeMap<String, FieldOrigin>,
}
impl Collector {
fn add(&mut self, field: &str, rule_title: &str, source: FieldSource) {
let entry = self.fields.entry(field.to_string()).or_default();
entry.rule_titles.insert(rule_title.to_string());
entry.sources.insert(source);
}
fn collect_detection_items(
&mut self,
detection: &Detection,
rule_title: &str,
source: FieldSource,
) {
match detection {
Detection::AllOf(items) => {
for item in items {
self.collect_item(item, rule_title, source);
}
}
Detection::AnyOf(subs) => {
for sub in subs {
self.collect_detection_items(sub, rule_title, source);
}
}
Detection::ArrayMatch { field, body, .. } => {
self.add(field, rule_title, source);
self.collect_detection_items(body, rule_title, source);
}
Detection::And(subs) => {
for sub in subs {
self.collect_detection_items(sub, rule_title, source);
}
}
Detection::Conditional { named, .. } => {
for sub in named.values() {
self.collect_detection_items(sub, rule_title, source);
}
}
Detection::Keywords(_) => {}
}
}
fn collect_item(&mut self, item: &DetectionItem, rule_title: &str, source: FieldSource) {
if let Some(ref name) = item.field.name {
self.add(name, rule_title, source);
}
}
fn collect_detections(
&mut self,
detections: &Detections,
rule_title: &str,
source: FieldSource,
) {
for det in detections.named.values() {
self.collect_detection_items(det, rule_title, source);
}
}
fn collect_rule(&mut self, rule: &SigmaRule) {
self.collect_detections(&rule.detection, &rule.title, FieldSource::Detection);
for f in &rule.fields {
self.add(f, &rule.title, FieldSource::Metadata);
}
}
fn collect_correlation(&mut self, corr: &CorrelationRule) {
for f in &corr.group_by {
self.add(f, &corr.title, FieldSource::Correlation);
}
if let CorrelationCondition::Threshold {
field: Some(ref fields),
..
} = corr.condition
{
for f in fields {
self.add(f, &corr.title, FieldSource::Correlation);
}
}
for alias in &corr.aliases {
for mapped_field in alias.mapping.values() {
self.add(mapped_field, &corr.title, FieldSource::Correlation);
}
}
for f in &corr.fields {
self.add(f, &corr.title, FieldSource::Metadata);
}
}
fn collect_filter(&mut self, filter: &FilterRule) {
self.collect_detections(&filter.detection, &filter.title, FieldSource::Filter);
for f in &filter.fields {
self.add(f, &filter.title, FieldSource::Metadata);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use rsigma_parser::parse_sigma_yaml;
fn build(yaml: &str) -> SigmaCollection {
parse_sigma_yaml(yaml).expect("parse")
}
#[test]
fn collects_detection_fields() {
let collection = build(
r#"
title: Test
status: test
logsource:
category: test
detection:
selection:
CommandLine|contains: whoami
EventID: 1
condition: selection
"#,
);
let set = RuleFieldSet::collect(&collection, &[], true);
assert!(set.contains("CommandLine"));
assert!(set.contains("EventID"));
assert!(
set.origin("CommandLine")
.unwrap()
.sources
.contains(&FieldSource::Detection)
);
}
#[test]
fn collects_correlation_group_by() {
let collection = build(
r#"
title: Login
id: login-rule
logsource:
category: auth
detection:
selection:
EventType: login
condition: selection
---
title: Many Logins
correlation:
type: event_count
rules:
- login-rule
group-by:
- User
timespan: 60s
condition:
gte: 3
"#,
);
let set = RuleFieldSet::collect(&collection, &[], true);
assert!(set.contains("EventType"));
assert!(set.contains("User"));
let user_origin = set.origin("User").unwrap();
assert!(user_origin.sources.contains(&FieldSource::Correlation));
}
#[test]
fn include_filters_toggle() {
let collection = build(
r#"
title: Detection
status: test
logsource:
category: test
detection:
selection:
DetField: x
condition: selection
---
title: Filter
filter:
rules:
- non-existent
selection:
FilterField: y
condition: selection
"#,
);
let with_filters = RuleFieldSet::collect(&collection, &[], true);
let without_filters = RuleFieldSet::collect(&collection, &[], false);
assert!(with_filters.contains("FilterField"));
assert!(!without_filters.contains("FilterField"));
assert!(with_filters.contains("DetField"));
assert!(without_filters.contains("DetField"));
}
#[test]
fn empty_collection_is_empty_set() {
let collection = SigmaCollection::default();
let set = RuleFieldSet::collect(&collection, &[], true);
assert!(set.is_empty());
assert_eq!(set.len(), 0);
}
}