usiem/components/dataset/
ip_set.rs

1use crate::prelude::SiemIp;
2use crossbeam_channel::Sender;
3use serde::Serialize;
4use std::collections::BTreeSet;
5use std::sync::Arc;
6
7#[derive(Serialize, Debug)]
8pub enum UpdateIpSet {
9    Add(SiemIp),
10    Remove(SiemIp),
11    Replace(IpSetDataset),
12}
13#[derive(Debug, Clone)]
14pub struct IpSetSynDataset {
15    dataset: Arc<IpSetDataset>,
16    comm: Sender<UpdateIpSet>,
17}
18impl IpSetSynDataset {
19    pub fn new(dataset: Arc<IpSetDataset>, comm: Sender<UpdateIpSet>) -> Self {
20        Self { dataset, comm }
21    }
22    pub fn empty() -> Self {
23        let (sender, _) = crossbeam_channel::bounded(1);
24        Self {
25            dataset: Arc::new(IpSetDataset::new()),
26            comm: sender,
27        }
28    }
29    /// Used to add IP with custom information like tags.
30    pub fn insert(&self, ip: SiemIp) {
31        // Todo: improve with local cache to send retries
32        let _ = self.comm.try_send(UpdateIpSet::Add(ip));
33    }
34    pub fn remove(&self, ip: SiemIp) {
35        // Todo: improve with local cache to send retries
36        let _ = self.comm.try_send(UpdateIpSet::Remove(ip));
37    }
38    pub fn update(&self, data: IpSetDataset) {
39        // Todo: improve with local cache to send retries
40        let _ = self.comm.try_send(UpdateIpSet::Replace(data));
41    }
42    pub fn contains(&self, ip: &SiemIp) -> bool {
43        // Todo improve with cached content
44        self.dataset.contains(ip)
45    }
46}
47#[derive(Serialize, Debug)]
48pub struct IpSetDataset {
49    data4: BTreeSet<u32>,
50    data6: BTreeSet<u128>,
51}
52
53impl Default for IpSetDataset {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59impl IpSetDataset {
60    pub fn new() -> Self {
61        Self {
62            data4: BTreeSet::new(),
63            data6: BTreeSet::new(),
64        }
65    }
66    pub fn insert(&mut self, ip: SiemIp) {
67        match ip {
68            SiemIp::V4(ip) => {
69                self.data4.insert(ip);
70            }
71            SiemIp::V6(ip) => {
72                self.data6.insert(ip);
73            }
74        }
75    }
76    pub fn contains(&self, ip: &SiemIp) -> bool {
77        match ip {
78            SiemIp::V4(ip) => self.data4.contains(ip),
79            SiemIp::V6(ip) => self.data6.contains(ip),
80        }
81    }
82    pub fn internal_ref(&self) -> (&BTreeSet<u32>, &BTreeSet<u128>) {
83        (&self.data4, &self.data6)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    #[test]
91    fn should_be_in_set() {
92        let mut dataset = IpSetDataset::new();
93        dataset.insert(SiemIp::from_ip_str("192.168.1.1").unwrap());
94        assert_eq!(
95            dataset.contains(&SiemIp::from_ip_str("192.168.1.1").unwrap()),
96            true
97        );
98    }
99}
usiem/components/dataset/ip_set.rs

usiem/components/dataset/
ip_set.rs
