use crate::decision::Decision; use crate::error::Error; use crate::error::Result; use crate::rule::PatternToken; use crate::rule::PrefixPattern; use crate::rule::PrefixRule; use crate::rule::RuleMatch; use crate::rule::RuleRef; use multimap::MultiMap; use serde::Deserialize; use serde::Serialize; use std::sync::Arc; type HeuristicsFallback<'a> = Option<&'a dyn Fn(&[String]) -> Decision>; #[derive(Clone, Debug)] pub struct Policy { rules_by_program: MultiMap, } impl Policy { pub fn new(rules_by_program: MultiMap) -> Self { Self { rules_by_program } } pub fn empty() -> Self { Self::new(MultiMap::new()) } pub fn rules(&self) -> &MultiMap { &self.rules_by_program } pub fn add_prefix_rule(&mut self, prefix: &[String], decision: Decision) -> Result<()> { let (first_token, rest) = prefix .split_first() .ok_or_else(|| Error::InvalidPattern("prefix cannot be empty".to_string()))?; let rule: RuleRef = Arc::new(PrefixRule { pattern: PrefixPattern { first: Arc::from(first_token.as_str()), rest: rest .iter() .map(|token| PatternToken::Single(token.clone())) .collect::>() .into(), }, decision, }); self.rules_by_program.insert(first_token.clone(), rule); Ok(()) } pub fn check(&self, cmd: &[String], heuristics_fallback: &F) -> Evaluation where F: Fn(&[String]) -> Decision, { let matched_rules = self.matches_for_command(cmd, Some(heuristics_fallback)); Evaluation::from_matches(matched_rules) } pub fn check_multiple( &self, commands: Commands, heuristics_fallback: &F, ) -> Evaluation where Commands: IntoIterator, Commands::Item: AsRef<[String]>, F: Fn(&[String]) -> Decision, { let matched_rules: Vec = commands .into_iter() .flat_map(|command| { self.matches_for_command(command.as_ref(), Some(heuristics_fallback)) }) .collect(); Evaluation::from_matches(matched_rules) } pub fn matches_for_command( &self, cmd: &[String], heuristics_fallback: HeuristicsFallback<'_>, ) -> Vec { let mut matched_rules: Vec = match cmd.first() { Some(first) => self .rules_by_program .get_vec(first) .map(|rules| rules.iter().filter_map(|rule| rule.matches(cmd)).collect()) .unwrap_or_default(), None => Vec::new(), }; if let (true, Some(heuristics_fallback)) = (matched_rules.is_empty(), heuristics_fallback) { matched_rules.push(RuleMatch::HeuristicsRuleMatch { command: cmd.to_vec(), decision: heuristics_fallback(cmd), }); } matched_rules } } #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct Evaluation { pub decision: Decision, #[serde(rename = "matchedRules")] pub matched_rules: Vec, } impl Evaluation { pub fn is_match(&self) -> bool { self.matched_rules .iter() .any(|rule_match| !matches!(rule_match, RuleMatch::HeuristicsRuleMatch { .. })) } fn from_matches(matched_rules: Vec) -> Self { let decision = matched_rules .iter() .map(RuleMatch::decision) .max() .unwrap_or(Decision::Allow); Self { decision, matched_rules, } } }