From 77c5d7e202f4429f5f4c4196f9ff1129346c9068 Mon Sep 17 00:00:00 2001
From: Alex Kornitzer <a.kornitzer@me.com>
Date: Wed, 19 Oct 2022 15:09:26 +0100
Subject: [PATCH] feat: add support for mapping extensions

This allows for overriding of poor rules as discussed in #107.
---
 src/hunt.rs       | 58 +++++++++++++++++++++++++++++++++++++++++++++++
 src/rule/sigma.rs | 34 ++++++++++++++++++++++++++-
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/src/hunt.rs b/src/hunt.rs
index 5753cae..1bc573f 100644
--- a/src/hunt.rs
+++ b/src/hunt.rs
@@ -28,6 +28,20 @@ use crate::rule::{
 };
 use crate::value::Value;
 
+#[derive(Clone, Deserialize)]
+pub struct Precondition {
+    #[serde(rename = "for")]
+    for_: HashMap<String, String>,
+    #[serde(deserialize_with = "crate::ext::tau::deserialize_expression")]
+    pub filter: Expression,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct Extensions {
+    #[serde(default)]
+    preconditions: Option<Vec<Precondition>>,
+}
+
 #[derive(Clone, Deserialize)]
 pub struct Group {
     #[serde(skip, default = "Uuid::new_v4")]
@@ -43,6 +57,8 @@ pub struct Group {
 pub struct Mapping {
     #[serde(default)]
     pub exclusions: HashSet<String>,
+    #[serde(default)]
+    pub extensions: Option<Extensions>,
     pub groups: Vec<Group>,
     pub kind: FileKind,
     pub name: String,
@@ -154,6 +170,40 @@ impl HunterBuilder {
                 if let RuleKind::Chainsaw = mapping.rules {
                     anyhow::bail!("Chainsaw rules do not support mappings");
                 }
+                let mut preconds = HashMap::new();
+                if let Some(extensions) = &mapping.extensions {
+                    if let Some(preconditions) = &extensions.preconditions {
+                        for precondition in preconditions {
+                            for (rid, rule) in &rules {
+                                if let Rule::Sigma(sigma) = rule {
+                                    // FIXME: How do we handle multiple matches, for now we just take
+                                    // the latest, we chould probably just combine them into an AND?
+                                    if precondition.for_.is_empty() {
+                                        continue;
+                                    }
+                                    let mut matched = true;
+                                    for (f, v) in &precondition.for_ {
+                                        match sigma.find(&f) {
+                                            Some(value) => {
+                                                if value.as_str() != Some(v.as_str()) {
+                                                    matched = false;
+                                                    break;
+                                                }
+                                            }
+                                            None => {
+                                                matched = false;
+                                                break;
+                                            }
+                                        }
+                                    }
+                                    if matched {
+                                        preconds.insert(*rid, precondition.filter.clone());
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
                 mapping.groups.sort_by(|x, y| x.name.cmp(&y.name));
                 for group in mapping.groups {
                     let mapper = Mapper::from(group.fields);
@@ -165,6 +215,7 @@ impl HunterBuilder {
                             exclusions: mapping.exclusions.clone(),
                             filter: group.filter,
                             kind: mapping.rules.clone(),
+                            preconditions: preconds.clone(),
                         },
                         timestamp: group.timestamp,
 
@@ -240,6 +291,7 @@ pub enum HuntKind {
         exclusions: HashSet<String>,
         filter: Expression,
         kind: RuleKind,
+        preconditions: HashMap<Uuid, Expression>,
     },
     Rule {
         aggregate: Option<Aggregate>,
@@ -503,6 +555,7 @@ impl Hunter {
                         exclusions,
                         filter,
                         kind,
+                        preconditions,
                     } => {
                         if tau_engine::core::solve(filter, &mapped) {
                             let matches = &self
@@ -528,6 +581,11 @@ impl Hunter {
                                     if exclusions.contains(rule.name()) {
                                         return None;
                                     }
+                                    if let Some(filter) = preconditions.get(rid) {
+                                        if !tau_engine::core::solve(filter, &mapped) {
+                                            return None;
+                                        }
+                                    }
                                     if rule.solve(&mapped) {
                                         Some((*rid, rule))
                                     } else {
diff --git a/src/rule/sigma.rs b/src/rule/sigma.rs
index a7ef4ff..f31d1aa 100644
--- a/src/rule/sigma.rs
+++ b/src/rule/sigma.rs
@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};
 use std::fs::File;
 use std::io::prelude::*;
@@ -7,7 +8,7 @@ use anyhow::Result;
 use regex::Regex;
 use serde::{Deserialize, Serialize};
 use serde_yaml::{Mapping, Sequence, Value as Yaml};
-use tau_engine::Rule as Tau;
+use tau_engine::{Document, Rule as Tau};
 
 use super::{Level, Status};
 
@@ -39,6 +40,37 @@ pub struct Rule {
     pub tags: Option<Vec<String>>,
 }
 
+impl Document for Rule {
+    fn find(&self, key: &str) -> Option<tau_engine::Value> {
+        use tau_engine::Value as Tau;
+        // NOTE: We have not implemented all fields here...
+        match key {
+            "title" => Some(Tau::String(Cow::Borrowed(&self.name))),
+            "level" => Some(Tau::String(Cow::Owned(self.level.to_string()))),
+            "status" => Some(Tau::String(Cow::Owned(self.status.to_string()))),
+            "id" => self.id.as_ref().map(|id| Tau::String(Cow::Borrowed(&id))),
+            "logsource.category" => self
+                .logsource
+                .as_ref()
+                .and_then(|ls| ls.category.as_ref().map(|c| Tau::String(Cow::Borrowed(&c)))),
+            "logsource.definition" => self.logsource.as_ref().and_then(|ls| {
+                ls.definition
+                    .as_ref()
+                    .map(|c| Tau::String(Cow::Borrowed(&c)))
+            }),
+            "logsource.product" => self
+                .logsource
+                .as_ref()
+                .and_then(|ls| ls.product.as_ref().map(|c| Tau::String(Cow::Borrowed(&c)))),
+            "logsource.service" => self
+                .logsource
+                .as_ref()
+                .and_then(|ls| ls.service.as_ref().map(|c| Tau::String(Cow::Borrowed(&c)))),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct Aggregate {
     pub count: String,