support regex matching against specific keys

Signed-off-by: Jean Mertz <[email protected]>

Author: JeanMertz

aw-query/src/datatype.rs

@@ -4,7 +4,7 @@ use std::fmt;
 use super::functions;
 use super::QueryError;
 use aw_models::Event;
-use aw_transform::classify::{RegexRule, Rule};
+use aw_transform::classify::{KeyValueRule, RegexRule, Rule};
 
 use serde::{Serialize, Serializer};
 use serde_json::value::Value;
@@ -297,50 +297,74 @@ impl TryFrom<&DataType> for Rule {
                 ))
             }
         };
-        if rtype == "none" {
-            Ok(Self::None)
-        } else if rtype == "regex" {
-            let regex_val = match obj.get("regex") {
-                Some(regex_val) => regex_val,
-                None => {
-                    return Err(QueryError::InvalidFunctionParameters(
-                        "regex rule is missing the 'regex' field".to_string(),
-                    ))
-                }
-            };
-            let regex_str = match regex_val {
-                DataType::String(s) => s,
-                _ => {
-                    return Err(QueryError::InvalidFunctionParameters(
-                        "the regex field of the regex rule is not a string".to_string(),
-                    ))
-                }
-            };
-            let ignore_case_val = match obj.get("ignore_case") {
-                Some(case_val) => case_val,
-                None => &DataType::Bool(false),
-            };
-            let ignore_case = match ignore_case_val {
-                DataType::Bool(b) => b,
-                _ => {
+
+        match rtype.as_str() {
+            "none" => Ok(Self::None),
+            "regex" => parse_regex_rule(obj),
+            "keyvalue" => {
+                let Some(rules) = obj.get("rules") else {
                     return Err(QueryError::InvalidFunctionParameters(
-                        "the ignore_case field of the regex rule is not a bool".to_string(),
-                    ))
-                }
-            };
-            let regex_rule = match RegexRule::new(regex_str, *ignore_case) {
-                Ok(regex_rule) => regex_rule,
-                Err(err) => {
-                    return Err(QueryError::RegexCompileError(format!(
-                        "Failed to compile regex string '{regex_str}': '{err:?}"
-                    )))
-                }
-            };
-            Ok(Self::Regex(regex_rule))
-        } else {
-            Err(QueryError::InvalidFunctionParameters(format!(
+                        "keyval rule is missing the 'rules' field".to_string(),
+                    ));
+                };
+
+                let rules = match rules {
+                    DataType::Dict(rules) => rules
+                        .iter()
+                        .map(|(k, v)| Rule::try_from(v).map(|v| (k.to_owned(), v)))
+                        .collect::<Result<HashMap<_, _>, _>>()?,
+                    _ => {
+                        return Err(QueryError::InvalidFunctionParameters(
+                            "the rules field of the keyval rule is not a dict".to_string(),
+                        ))
+                    }
+                };
+
+                Ok(Rule::KeyValue(KeyValueRule::new(rules)))
+            }
+            _ => Err(QueryError::InvalidFunctionParameters(format!(
                 "Unknown rule type '{rtype}'"
-            )))
+            ))),
         }
     }
 }
+
+fn parse_regex_rule(obj: &HashMap<String, DataType>) -> Result<Rule, QueryError> {
+    let regex_val = match obj.get("regex") {
+        Some(regex_val) => regex_val,
+        None => {
+            return Err(QueryError::InvalidFunctionParameters(
+                "regex rule is missing the 'regex' field".to_string(),
+            ))
+        }
+    };
+    let regex_str = match regex_val {
+        DataType::String(s) => s,
+        _ => {
+            return Err(QueryError::InvalidFunctionParameters(
+                "the regex field of the regex rule is not a string".to_string(),
+            ))
+        }
+    };
+    let ignore_case_val = match obj.get("ignore_case") {
+        Some(case_val) => case_val,
+        None => &DataType::Bool(false),
+    };
+    let ignore_case = match ignore_case_val {
+        DataType::Bool(b) => b,
+        _ => {
+            return Err(QueryError::InvalidFunctionParameters(
+                "the ignore_case field of the regex rule is not a bool".to_string(),
+            ))
+        }
+    };
+    let regex_rule = match RegexRule::new(regex_str, *ignore_case) {
+        Ok(regex_rule) => regex_rule,
+        Err(err) => {
+            return Err(QueryError::RegexCompileError(format!(
+                "Failed to compile regex string '{regex_str}': '{err:?}"
+            )))
+        }
+    };
+    Ok(Rule::Regex(regex_rule))
+}

aw-transform/src/classify.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 /// Transforms for classifying (tagging and categorizing) events.
 ///
 /// Based on code in aw_research: https://github.com/ActivityWatch/aw-research/blob/master/aw_research/classify.py
@@ -7,13 +9,15 @@ use fancy_regex::Regex;
 pub enum Rule {
     None,
     Regex(RegexRule),
+    KeyValue(KeyValueRule),
 }
 
 impl RuleTrait for Rule {
     fn matches(&self, event: &Event) -> bool {
         match self {
             Rule::None => false,
             Rule::Regex(rule) => rule.matches(event),
+            Rule::KeyValue(rule) => rule.matches(event),
         }
     }
 }
@@ -62,6 +66,32 @@ impl From<Regex> for Rule {
     }
 }
 
+pub struct KeyValueRule {
+    rules: HashMap<String, Rule>,
+}
+
+impl KeyValueRule {
+    pub fn new(rules: HashMap<String, Rule>) -> Self {
+        Self { rules }
+    }
+}
+
+impl RuleTrait for KeyValueRule {
+    fn matches(&self, event: &Event) -> bool {
+        self.rules.iter().all(|(key, rule)| {
+            event
+                .data
+                .get(key)
+                .filter(|_| {
+                    let mut ev = event.clone();
+                    ev.data.retain(|k, _| k == key);
+                    rule.matches(&ev)
+                })
+                .is_some()
+        })
+    }
+}
+
 /// Categorizes a list of events
 ///
 /// An event can only have one category, although the category may have a hierarchy,