use crate::decision::Decision; use crate::error::Error; use crate::error::Result; use crate::rule::NetworkRule; use crate::rule::NetworkRuleProtocol; use crate::rule::PatternToken; use crate::rule::PrefixPattern; use crate::rule::PrefixRule; use crate::rule::RuleMatch; use crate::rule::RuleRef; use crate::rule::normalize_network_rule_host; 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, network_rules: Vec, } impl Policy { pub fn new(rules_by_program: MultiMap) -> Self { Self::from_parts(rules_by_program, Vec::new()) } pub fn from_parts( rules_by_program: MultiMap, network_rules: Vec, ) -> Self { Self { rules_by_program, network_rules, } } pub fn empty() -> Self { Self::new(MultiMap::new()) } pub fn rules(&self) -> &MultiMap { &self.rules_by_program } pub fn network_rules(&self) -> &[NetworkRule] { &self.network_rules } pub fn get_allowed_prefixes(&self) -> Vec> { let mut prefixes = Vec::new(); for (_program, rules) in self.rules_by_program.iter_all() { for rule in rules { let Some(prefix_rule) = rule.as_any().downcast_ref::() else { continue; }; if prefix_rule.decision != Decision::Allow { continue; } let mut prefix = Vec::with_capacity(prefix_rule.pattern.rest.len() + 1); prefix.push(prefix_rule.pattern.first.as_ref().to_string()); prefix.extend(prefix_rule.pattern.rest.iter().map(render_pattern_token)); prefixes.push(prefix); } } prefixes.sort(); prefixes.dedup(); prefixes } 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, justification: None, }); self.rules_by_program.insert(first_token.clone(), rule); Ok(()) } pub fn add_network_rule( &mut self, host: &str, protocol: NetworkRuleProtocol, decision: Decision, justification: Option, ) -> Result<()> { let host = normalize_network_rule_host(host)?; if let Some(raw) = justification.as_deref() && raw.trim().is_empty() { return Err(Error::InvalidRule( "justification cannot be empty".to_string(), )); } self.network_rules.push(NetworkRule { host, protocol, decision, justification, }); Ok(()) } pub fn compiled_network_domains(&self) -> (Vec, Vec) { let mut allowed = Vec::new(); let mut denied = Vec::new(); for rule in &self.network_rules { match rule.decision { Decision::Allow => { denied.retain(|entry| entry != &rule.host); upsert_domain(&mut allowed, &rule.host); } Decision::Forbidden => { allowed.retain(|entry| entry != &rule.host); upsert_domain(&mut denied, &rule.host); } Decision::Prompt => {} } } (allowed, denied) } 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) } /// Checks multiple commands and aggregates the results. 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) } /// Returns matching rules for the given command. If no rules match and /// `heuristics_fallback` is provided, returns a single /// `HeuristicsRuleMatch` with the decision rendered by /// `heuristics_fallback`. /// /// If `heuristics_fallback.is_some()`, then the returned vector is /// guaranteed to be non-empty. pub fn matches_for_command( &self, cmd: &[String], heuristics_fallback: HeuristicsFallback<'_>, ) -> Vec { let 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 matched_rules.is_empty() && let Some(heuristics_fallback) = heuristics_fallback { vec![RuleMatch::HeuristicsRuleMatch { command: cmd.to_vec(), decision: heuristics_fallback(cmd), }] } else { matched_rules } } } fn upsert_domain(entries: &mut Vec, host: &str) { entries.retain(|entry| entry != host); entries.push(host.to_string()); } fn render_pattern_token(token: &PatternToken) -> String { match token { PatternToken::Single(value) => value.clone(), PatternToken::Alts(alternatives) => format!("[{}]", alternatives.join("|")), } } #[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 { .. })) } /// Caller is responsible for ensuring that `matched_rules` is non-empty. fn from_matches(matched_rules: Vec) -> Self { let decision = matched_rules.iter().map(RuleMatch::decision).max(); #[expect(clippy::expect_used)] let decision = decision.expect("invariant failed: matched_rules must be non-empty"); Self { decision, matched_rules, } } }