usiem/components/dataset/
ip_map.rs

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

usiem/components/dataset/
ip_map.rs
