usiem/components/dataset/
rules.rs

1use crate::prelude::rule::SiemRule;
2use crate::prelude::types::LogString;
3use crossbeam_channel::Sender;
4use serde::Serialize;
5use std::collections::BTreeMap;
6use std::sync::Arc;
7use std::time::Duration;
8
9#[derive(Serialize, Debug)]
10pub enum UpdateRules {
11    Add(SiemRule),
12    Remove(LogString),
13    Replace(RulesDataset),
14}
15
16#[derive(Debug, Clone)]
17pub struct CorrelationRulesDataset {
18    dataset: Arc<RulesDataset>,
19    comm: Sender<UpdateRules>,
20}
21impl CorrelationRulesDataset {
22    pub fn new(dataset: Arc<RulesDataset>, comm: Sender<UpdateRules>) -> Self {
23        Self { dataset, comm }
24    }
25    pub fn insert(&self, rule: SiemRule) {
26        let _ = self.comm.send(UpdateRules::Add(rule));
27    }
28    pub fn insert_timeout(&self, rule: SiemRule, timeout: Duration) -> Result<(), SiemRule> {
29        let init = std::time::Instant::now();
30        let mut rule = rule;
31        loop {
32            rule = match self.comm.try_send(UpdateRules::Add(rule)) {
33                Ok(_) => return Ok(()),
34                Err(e) => match e {
35                    crossbeam_channel::TrySendError::Full(r) => extract_rule_from_update(r),
36                    crossbeam_channel::TrySendError::Disconnected(r) => extract_rule_from_update(r),
37                },
38            };
39            let now = std::time::Instant::now();
40            if now > init + timeout {
41                return Err(rule);
42            }
43        }
44    }
45    pub fn try_insert(&self, rule: SiemRule) -> Result<(), SiemRule> {
46        match self.comm.try_send(UpdateRules::Add(rule)) {
47            Ok(_) => Ok(()),
48            Err(e) => match e {
49                crossbeam_channel::TrySendError::Full(r) => Err(extract_rule_from_update(r)),
50                crossbeam_channel::TrySendError::Disconnected(r) => {
51                    Err(extract_rule_from_update(r))
52                }
53            },
54        }
55    }
56    pub fn get(&self, id: &LogString) -> Option<&SiemRule> {
57        // Todo improve with cached added IPs
58        self.dataset.get(id)
59    }
60}
61#[derive(Serialize, Debug, Default)]
62pub struct RulesDataset {
63    rules: BTreeMap<LogString, SiemRule>,
64}
65
66impl RulesDataset {
67    pub fn new() -> Self {
68        Self::default()
69    }
70    pub fn insert(&mut self, rule: SiemRule) {
71        self.rules.insert(rule.name.clone(), rule);
72    }
73    pub fn get(&self, id: &LogString) -> Option<&SiemRule> {
74        self.rules.get(id)
75    }
76}
77
78fn extract_rule_from_update(update: UpdateRules) -> SiemRule {
79    match update {
80        UpdateRules::Add(r) => r,
81        UpdateRules::Remove(_) => unreachable!(),
82        UpdateRules::Replace(_) => unreachable!(),
83    }
84}
usiem/components/dataset/rules.rs

usiem/components/dataset/
rules.rs
