diff --git a/src/config/mod.rs b/src/config/mod.rs index 18be813..c9c1c14 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -167,44 +167,34 @@ impl Context { res } - pub async fn update_blocklist(&mut self, ipdata: &IpData) -> Option { - if self.blocklist.contains_key(&ipdata.ip) - && self.hostname == ipdata.hostname - && ipdata.mode == "file".to_string() - { - let mut block = self.blocklist.get_mut(&ipdata.ip).unwrap(); - block.tryfail += 1; - match self.cfg.sets.get(&ipdata.src) { - Some(set) => { + pub async fn update_blocklist(&mut self, ipdata: &mut IpData) -> Option { + match self.cfg.sets.get(&ipdata.src) { + Some(set) => { + if self.blocklist.contains_key(&ipdata.ip) + && self.hostname == ipdata.hostname + && ipdata.mode == "file".to_string() + { + let mut block = self.blocklist.get_mut(&ipdata.ip).unwrap(); + block.tryfail += 1; + block.blocktime = set.blocktime; if block.tryfail >= set.tryfail { return Some(block.ipdata.clone()); } - } - None => {} - } - } - if !self.blocklist.contains_key(&ipdata.ip) { - let mut tryfail = 0; - match self.cfg.sets.get(&ipdata.src) { - Some(set) => { + } else { let starttime: DateTime = DateTime::parse_from_rfc3339(ipdata.date.as_str()).unwrap(); - if ipdata.mode == "zmq".to_string() { - tryfail = 100; - } - self.blocklist.insert( - ipdata.ip.to_string(), - BlockIpData { + self.blocklist + .entry(ipdata.ip.to_string()) + .or_insert(BlockIpData { ipdata: ipdata.clone(), - tryfail, + tryfail: 100, starttime, blocktime: set.blocktime, - }, - ); + }); } - None => {} } - }; + None => {} + } None } @@ -462,7 +452,7 @@ mod test { ctx.blocklist = HashMap::new(); for _i in 0..10 { - ctx.update_blocklist(&IpData { + ctx.update_blocklist(&mut IpData { ip: "1.1.1.1".to_string(), hostname: "test1".to_string(), date: now.to_rfc3339().to_string(), @@ -473,7 +463,7 @@ mod test { } for _i in 0..10 { - ctx.update_blocklist(&IpData { + ctx.update_blocklist(&mut IpData { ip: "1.1.1.2".to_string(), hostname: "test2".to_string(), date: now.to_rfc3339().to_string(), @@ -483,7 +473,7 @@ mod test { .await; } - ctx.update_blocklist(&IpData { + ctx.update_blocklist(&mut IpData { ip: "1.1.1.3".to_string(), hostname: "testgood".to_string(), date: now.to_rfc3339().to_string(), @@ -492,7 +482,7 @@ mod test { }) .await; - ctx.update_blocklist(&IpData { + ctx.update_blocklist(&mut IpData { ip: "1.1.1.4".to_string(), hostname: "testgood".to_string(), date: now.to_rfc3339().to_string(), @@ -501,7 +491,7 @@ mod test { }) .await; - ctx.update_blocklist(&IpData { + ctx.update_blocklist(&mut IpData { ip: "1.1.1.4".to_string(), hostname: "testgood".to_string(), date: now.to_rfc3339().to_string(), diff --git a/src/ipblc/inc.rs b/src/ipblc/inc.rs index a4c8552..fa90575 100644 --- a/src/ipblc/inc.rs +++ b/src/ipblc/inc.rs @@ -60,22 +60,22 @@ pub async fn process(ctx: &Arc>) { let begin: DateTime = Local::now().trunc_subsecs(0); // wait for logs parse and zmq channel receive - let received_ip = ipdatarx.recv().await.unwrap(); + let mut received_ip = ipdatarx.recv().await.unwrap(); // lock the context mutex let ctxarc = Arc::clone(&ctx); let mut ctx = ctxarc.lock().await; - if received_ip.mode == "init" { + if received_ip.ip == "".to_string() && received_ip.mode == "init".to_string() { for ip_to_send in &mut ctx.get_blocklist_toblock().await { - ip_to_send.mode = "zmq".to_string(); + ip_to_send.mode = "init".to_string(); send_to_ipbl_zmq(&reqsocket, ip_to_send).await; } continue; } // refresh context blocklist - let filtered_ip = ctx.update_blocklist(&received_ip).await; + let filtered_ip = ctx.update_blocklist(&mut received_ip).await; ctx.gc_blocklist().await; // send ip list to ws and zmq sockets diff --git a/src/ipblc/mod.rs b/src/ipblc/mod.rs index 1b9a2e2..ea49901 100644 --- a/src/ipblc/mod.rs +++ b/src/ipblc/mod.rs @@ -25,7 +25,6 @@ impl std::fmt::Debug for FileEvent { async fn send_to_ipbl_zmq(reqsocket: &zmq::Socket, ip: &mut IpData) { let msg = format!("{value}", value = serde_json::to_string(&ip).unwrap()); - ip.mode = "zmq".to_string(); match reqsocket.send(&msg, 0) { Ok(_) => {} Err(e) => { @@ -47,7 +46,6 @@ async fn send_to_ipbl_zmq(reqsocket: &zmq::Socket, ip: &mut IpData) { async fn send_to_ipbl_ws(ctx: &Context, ip: &mut IpData, ret: &mut Vec) { ret.push(format!("host: {hostname}", hostname = ctx.hostname)); - ip.mode = "zmq".to_string(); loop { match push_ip(&ctx, &ip, ret).await { Ok(_) => { @@ -91,7 +89,7 @@ async fn listenpubsub(ctx: &Arc>, txpubsub: Sender, socke Some(ss) => { let msg = ss.strip_prefix(prefix.as_str()).unwrap(); let tosend: IpData = serde_json::from_str(msg).unwrap(); - if tosend.hostname != gethostname(true) || tosend.mode == "zmq".to_string() { + if tosend.hostname != gethostname(true) || tosend.mode == "init".to_string() { txpubsub.send(tosend).await.unwrap(); } }