diff --git a/src/config.rs b/src/config.rs index 0e17df9..3fae15b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,4 @@ -use crate::ip::{BlockIpData, IpData}; +use crate::ip::{BlockIpData, IpData, IpEvent}; use crate::utils::{gethostname, sleep_s}; use chrono::prelude::*; @@ -156,15 +156,11 @@ impl Context { pub async fn get_blocklist_toblock(&mut self) -> Vec { let mut res: Vec = vec![]; - //let now: DateTime = Local::now().trunc_subsecs(0); for (_, block) in self.blocklist.iter_mut() { match self.cfg.sets.get(&block.ipdata.src) { Some(set) => { if block.tryfail >= set.tryfail { res.push(block.ipdata.clone()); - /*if block.tryfail == set.tryfail { - block.starttime = DateTime::from(now); - }*/ } } None => {} @@ -173,27 +169,27 @@ impl Context { res } - pub async fn update_blocklist(&mut self, ipdata: &IpData) -> Option { - match self.cfg.sets.get(&ipdata.src) { + pub async fn update_blocklist(&mut self, ipevent: &IpEvent) -> Option { + match self.cfg.sets.get(&ipevent.ipdata.src) { Some(set) => { - if self.blocklist.contains_key(&ipdata.ip) - && self.hostname == ipdata.hostname - && ipdata.mode == "file".to_string() + if self.blocklist.contains_key(&ipevent.ipdata.ip) + && self.hostname == ipevent.ipdata.hostname + && ipevent.mode == "file".to_string() { - let mut block = self.blocklist.get_mut(&ipdata.ip).unwrap(); + let mut block = self.blocklist.get_mut(&ipevent.ipdata.ip).unwrap(); block.tryfail += 1; block.blocktime = set.blocktime; if block.tryfail >= set.tryfail { - return Some(block.ipdata.clone()); + return Some(ipevent.ipdata.clone()); } } else { - let starttime = DateTime::parse_from_rfc3339(ipdata.date.as_str()) + let starttime = DateTime::parse_from_rfc3339(ipevent.ipdata.date.as_str()) .unwrap() .with_timezone(&chrono::Local); self.blocklist - .entry(ipdata.ip.to_string()) + .entry(ipevent.ipdata.ip.to_string()) .or_insert(BlockIpData { - ipdata: ipdata.clone(), + ipdata: ipevent.ipdata.clone(), tryfail: set.tryfail, starttime, blocktime: set.blocktime, @@ -209,7 +205,7 @@ impl Context { let mut removed: Vec = vec![]; let now: DateTime = Local::now().trunc_subsecs(0); // nightly, future use - //let drained: HashMap = ctx.blocklist.drain_filter(|k,v| v.parse_date() < mindate) + // let drained: HashMap = ctx.blocklist.drain_filter(|k,v| v.parse_date() < mindate) for (ip, blocked) in self.blocklist.clone().iter() { match self.cfg.sets.get(&blocked.ipdata.src) { Some(set) => { @@ -530,51 +526,66 @@ mod test { ctx.blocklist = HashMap::new(); for _i in 0..10 { - ctx.update_blocklist(&mut IpData { - ip: "1.1.1.1".to_string(), - hostname: "test1".to_string(), - date: now.to_rfc3339().to_string(), - src: "ssh".to_string(), - mode: "file".to_string(), + ctx.update_blocklist(&mut IpEvent { + msgtype: String::from("add"), + mode: String::from("zmq"), + ipdata: IpData { + ip: "1.1.1.1".to_string(), + hostname: "test1".to_string(), + date: now.to_rfc3339().to_string(), + src: "ssh".to_string(), + }, }) .await; } for _ in 0..10 { - ctx.update_blocklist(&mut IpData { - ip: "1.1.1.2".to_string(), - hostname: "test2".to_string(), - date: now.to_rfc3339().to_string(), - src: "http".to_string(), - mode: "file".to_string(), + ctx.update_blocklist(&mut IpEvent { + msgtype: String::from("add"), + mode: String::from("zmq"), + ipdata: IpData { + ip: "1.1.1.2".to_string(), + hostname: "test2".to_string(), + date: now.to_rfc3339().to_string(), + src: "http".to_string(), + }, }) .await; } - ctx.update_blocklist(&mut IpData { - ip: "1.1.1.3".to_string(), - hostname: "testgood".to_string(), - date: now.to_rfc3339().to_string(), - src: "http".to_string(), - mode: "file".to_string(), + ctx.update_blocklist(&mut IpEvent { + msgtype: String::from("add"), + mode: String::from("zmq"), + ipdata: IpData { + ip: "1.1.1.3".to_string(), + hostname: "testgood".to_string(), + date: now.to_rfc3339().to_string(), + src: "http".to_string(), + }, }) .await; - ctx.update_blocklist(&mut IpData { - ip: "1.1.1.4".to_string(), - hostname: "testgood".to_string(), - date: now.to_rfc3339().to_string(), - src: "http".to_string(), - mode: "file".to_string(), + ctx.update_blocklist(&mut IpEvent { + msgtype: String::from("add"), + mode: String::from("zmq"), + ipdata: IpData { + ip: "1.1.1.4".to_string(), + hostname: "testgood".to_string(), + date: now.to_rfc3339().to_string(), + src: "http".to_string(), + }, }) .await; - ctx.update_blocklist(&mut IpData { - ip: "1.1.1.4".to_string(), - hostname: "testgood".to_string(), - date: now.to_rfc3339().to_string(), - src: "http".to_string(), - mode: "file".to_string(), + ctx.update_blocklist(&mut IpEvent { + msgtype: String::from("add"), + mode: String::from("zmq"), + ipdata: IpData { + ip: "1.1.1.4".to_string(), + hostname: "testgood".to_string(), + date: now.to_rfc3339().to_string(), + src: "http".to_string(), + }, }) .await; diff --git a/src/ip.rs b/src/ip.rs index b11de03..caeb47f 100644 --- a/src/ip.rs +++ b/src/ip.rs @@ -21,6 +21,7 @@ lazy_static! { #[derive(Clone, Debug, Serialize, Deserialize)] pub struct IpEvent { pub msgtype: String, + pub mode: String, pub ipdata: IpData, } @@ -30,7 +31,6 @@ pub struct IpData { pub src: String, pub date: String, pub hostname: String, - pub mode: String, } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -66,12 +66,11 @@ impl Display for IpData { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { write!( f, - "ip: {ip}, src: {src}, date: {date}, hostname: {hostname}, mode: {mode}", + "ip: {ip}, src: {src}, date: {date}, hostname: {hostname}", ip = self.ip, src = self.src, date = self.date, hostname = self.hostname, - mode = self.mode, ) } } @@ -134,7 +133,6 @@ pub fn filter( src: src.to_owned(), date: s_date.to_rfc3339().to_owned(), hostname: hostname.to_owned(), - mode: "file".to_owned(), }); ips += 1; }; diff --git a/src/ipblc.rs b/src/ipblc.rs index f135fb4..112d16c 100644 --- a/src/ipblc.rs +++ b/src/ipblc.rs @@ -53,12 +53,12 @@ pub async fn run() { let ipevent_bootstrap = IpEvent { msgtype: String::from("bootstrap"), + mode: String::from("zmq"), ipdata: IpData { ip: "".to_string(), src: "".to_string(), date: "".to_string(), hostname: "".to_string(), - mode: "".to_string(), }, }; send_to_ipbl_zmq(&zmqreqsocket, &ipevent_bootstrap, &mut ret).await; @@ -84,6 +84,7 @@ pub async fn run() { for ip_to_send in ctx.get_blocklist_toblock().await { let ipe = IpEvent{ msgtype: String::from("init"), + mode: String::from("zmq"), ipdata: ip_to_send, }; send_to_ipbl_zmq(&zmqreqsocket, &ipe, &mut ret).await; @@ -92,14 +93,16 @@ pub async fn run() { } // refresh context blocklist - let filtered_ip = ctx.update_blocklist(&received_ip.ipdata).await; + let filtered_ip = ctx.update_blocklist(&received_ip).await; // send ip list to ws and zmq sockets if let Some(ip) = filtered_ip { + println!("{}",ip); if received_ip.msgtype != "init" { println!("sending {} to ws and zmq", ip.ip); let event = IpEvent{ msgtype: String::from("add"), + mode:String::from("zmq"), ipdata: ip, }; send_to_ipbl_ws(&ctx, &event, &mut ret).await; @@ -248,6 +251,7 @@ async fn compare_files_changes( for ip in iplist { let ipevent = IpEvent { msgtype: String::from("add"), + mode: String::from("file"), ipdata: ip, }; ipeventtx.send(ipevent).await.unwrap(); diff --git a/src/ws.rs b/src/ws.rs index 20818a7..9d48fae 100644 --- a/src/ws.rs +++ b/src/ws.rs @@ -28,7 +28,6 @@ async fn push_ip(ctx: &Context, ip: &IpData, ret: &mut Vec) -> Result<() src: ip.src.to_string(), date: ip.date.to_string(), hostname: ip.hostname.to_string(), - mode: "file".to_string(), }); let resp = ctx @@ -65,7 +64,6 @@ async fn _push_ip_bulk( src: ip.src.to_string(), date: ip.date.to_string(), hostname: ip.hostname.to_string(), - mode: "file".to_string(), }) }