psmatcher 0.4.0-alpha.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{fmt, str::FromStr};

use regex::Regex;
use serde::{Deserialize, Serialize};

use crate::{
    error::MatcherError,
    event::{AttributeValue, Event},
};

/// Custom serialization for Regex objects
mod regex_serde {
    use std::str::FromStr;

    use regex::Regex;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S>(regex: &Regex, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        regex.as_str().serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Regex, D::Error>
    where
        D: Deserializer<'de>,
    {
        let pattern = String::deserialize(deserializer)?;
        Regex::from_str(&pattern).map_err(serde::de::Error::custom)
    }
}

/// Represents an elementary test that can be performed on an event attribute
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ElementaryTest {
    /// Tests if an attribute equals a value
    Equals {
        /// The name of the attribute to test
        attribute: String,
        /// The value to compare against
        value: AttributeValue,
    },

    /// Tests if an attribute contains a substring (for string attributes)
    Contains {
        /// The name of the attribute to test
        attribute: String,
        /// The substring to search for
        substring: String,
    },

    /// Tests if an attribute starts with a prefix (for string attributes)
    StartsWith {
        /// The name of the attribute to test
        attribute: String,
        /// The prefix to check
        prefix: String,
    },

    /// Tests if an attribute ends with a suffix (for string attributes)
    EndsWith {
        /// The name of the attribute to test
        attribute: String,
        /// The suffix to check
        suffix: String,
    },

    /// Tests if an attribute matches a regex pattern (for string attributes)
    Regex {
        /// The name of the attribute to test
        attribute: String,
        /// The compiled regex pattern to match
        #[serde(with = "regex_serde")]
        pattern: Regex,
    },

    /// Tests if a numeric attribute is greater than a value
    GreaterThan {
        /// The name of the attribute to test
        attribute: String,
        /// The value to compare against
        value: f64,
    },

    /// Tests if a numeric attribute is less than a value
    LessThan {
        /// The name of the attribute to test
        attribute: String,
        /// The value to compare against
        value: f64,
    },

    /// Tests if an attribute is in a list of values
    In {
        /// The name of the attribute to test
        attribute: String,
        /// The list of values to check against
        values: Vec<AttributeValue>,
    },

    /// Tests if an attribute exists
    Exists {
        /// The name of the attribute to check for existence
        attribute: String,
    },
}

// Implement PartialEq manually to handle Regex comparison
impl PartialEq for ElementaryTest {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (
                ElementaryTest::Equals {
                    attribute: a1,
                    value: v1,
                },
                ElementaryTest::Equals {
                    attribute: a2,
                    value: v2,
                },
            ) => a1 == a2 && v1 == v2,
            (
                ElementaryTest::Contains {
                    attribute: a1,
                    substring: s1,
                },
                ElementaryTest::Contains {
                    attribute: a2,
                    substring: s2,
                },
            ) => a1 == a2 && s1 == s2,
            (
                ElementaryTest::StartsWith {
                    attribute: a1,
                    prefix: p1,
                },
                ElementaryTest::StartsWith {
                    attribute: a2,
                    prefix: p2,
                },
            ) => a1 == a2 && p1 == p2,
            (
                ElementaryTest::EndsWith {
                    attribute: a1,
                    suffix: s1,
                },
                ElementaryTest::EndsWith {
                    attribute: a2,
                    suffix: s2,
                },
            ) => a1 == a2 && s1 == s2,
            (
                ElementaryTest::Regex {
                    attribute: a1,
                    pattern: p1,
                },
                ElementaryTest::Regex {
                    attribute: a2,
                    pattern: p2,
                },
            ) => a1 == a2 && p1.as_str() == p2.as_str(),
            (
                ElementaryTest::GreaterThan {
                    attribute: a1,
                    value: v1,
                },
                ElementaryTest::GreaterThan {
                    attribute: a2,
                    value: v2,
                },
            ) => a1 == a2 && v1 == v2,
            (
                ElementaryTest::LessThan {
                    attribute: a1,
                    value: v1,
                },
                ElementaryTest::LessThan {
                    attribute: a2,
                    value: v2,
                },
            ) => a1 == a2 && v1 == v2,
            (
                ElementaryTest::In {
                    attribute: a1,
                    values: v1,
                },
                ElementaryTest::In {
                    attribute: a2,
                    values: v2,
                },
            ) => a1 == a2 && v1 == v2,
            (
                ElementaryTest::Exists { attribute: a1 },
                ElementaryTest::Exists { attribute: a2 },
            ) => a1 == a2,
            _ => false,
        }
    }
}

impl ElementaryTest {
    /// Creates a new Regex test with the given attribute name and pattern
    pub fn regex(attribute: String, pattern: &str) -> Result<Self, MatcherError> {
        match Regex::from_str(pattern) {
            Ok(regex) => Ok(ElementaryTest::Regex {
                attribute,
                pattern: regex,
            }),
            Err(err) => Err(MatcherError::InvalidPredicate(format!(
                "Invalid regex pattern: {err}"
            ))),
        }
    }

    /// Evaluates the test against an event
    pub fn evaluate(&self, event: &impl Event) -> Result<bool, MatcherError> {
        match self {
            ElementaryTest::Equals { attribute, value } => match event.get_attribute(attribute) {
                Some(attr_value) => Ok(attr_value == value),
                None => Ok(false),
            },

            ElementaryTest::Contains {
                attribute,
                substring,
            } => match event.get_attribute(attribute) {
                Some(AttributeValue::String(s)) => Ok(s.contains(substring)),
                Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                None => Ok(false),
            },

            ElementaryTest::StartsWith { attribute, prefix } => {
                match event.get_attribute(attribute) {
                    Some(AttributeValue::String(s)) => Ok(s.starts_with(prefix)),
                    Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                    None => Ok(false),
                }
            }

            ElementaryTest::EndsWith { attribute, suffix } => {
                match event.get_attribute(attribute) {
                    Some(AttributeValue::String(s)) => Ok(s.ends_with(suffix)),
                    Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                    None => Ok(false),
                }
            }

            ElementaryTest::Regex { attribute, pattern } => match event.get_attribute(attribute) {
                Some(AttributeValue::String(s)) => Ok(pattern.is_match(s)),
                Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                None => Ok(false),
            },

            ElementaryTest::GreaterThan { attribute, value } => {
                match event.get_attribute(attribute) {
                    Some(AttributeValue::Integer(i)) => Ok(*i as f64 > *value),
                    Some(AttributeValue::Float(f)) => Ok(*f > *value),
                    Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                    None => Ok(false),
                }
            }

            ElementaryTest::LessThan { attribute, value } => match event.get_attribute(attribute) {
                Some(AttributeValue::Integer(i)) => Ok((*i as f64) < *value),
                Some(AttributeValue::Float(f)) => Ok(*f < *value),
                Some(_) => Err(MatcherError::IncompatibleAttributeType(attribute.clone())),
                None => Ok(false),
            },

            ElementaryTest::In { attribute, values } => match event.get_attribute(attribute) {
                Some(attr_value) => Ok(values.contains(attr_value)),
                None => Ok(false),
            },

            ElementaryTest::Exists { attribute } => Ok(event.get_attribute(attribute).is_some()),
        }
    }
}

/// Represents an elementary predicate, which is the result of an elementary test
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ElementaryPredicate {
    /// The test that this predicate represents
    pub test: ElementaryTest,
    /// Whether the result of the test should be negated
    pub negated: bool,
}

impl ElementaryPredicate {
    /// Creates a new elementary predicate
    pub fn new(test: ElementaryTest, negated: bool) -> Self {
        Self { test, negated }
    }

    /// Evaluates the predicate against an event
    pub fn evaluate(&self, event: &impl Event) -> Result<bool, MatcherError> {
        let result = self.test.evaluate(event)?;
        Ok(if self.negated { !result } else { result })
    }
}

// Implement PartialEq for ElementaryPredicate
impl PartialEq for ElementaryPredicate {
    fn eq(&self, other: &Self) -> bool {
        self.test == other.test && self.negated == other.negated
    }
}

impl fmt::Display for ElementaryPredicate {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let negation = if self.negated { "NOT " } else { "" };
        match &self.test {
            ElementaryTest::Equals { attribute, value } => {
                write!(f, "{negation}({attribute} = {value:?})")
            }
            ElementaryTest::Contains {
                attribute,
                substring,
            } => {
                write!(f, "{negation}({attribute} CONTAINS \"{substring}\")")
            }
            ElementaryTest::StartsWith { attribute, prefix } => {
                write!(f, "{negation}({attribute} STARTS WITH \"{prefix}\")")
            }
            ElementaryTest::EndsWith { attribute, suffix } => {
                write!(f, "{negation}({attribute} ENDS WITH \"{suffix}\")")
            }
            ElementaryTest::Regex { attribute, pattern } => {
                write!(
                    f,
                    "{}({} MATCHES \"{}\")",
                    negation,
                    attribute,
                    pattern.as_str()
                )
            }
            ElementaryTest::GreaterThan { attribute, value } => {
                write!(f, "{negation}({attribute} > {value})")
            }
            ElementaryTest::LessThan { attribute, value } => {
                write!(f, "{negation}({attribute} < {value})")
            }
            ElementaryTest::In { attribute, values } => {
                write!(f, "{negation}({attribute} IN {values:?})")
            }
            ElementaryTest::Exists { attribute } => {
                write!(f, "{negation}({attribute} EXISTS)")
            }
        }
    }
}