add grace period

src/jail.rs

@@ -5,53 +5,66 @@ use std::sync::Mutex;
 
 use anyhow::*;
 
-use crate::utils::log;
+use crate::utils::{get_epoch, log};
 
 pub struct Jail {
     name: String,
     allowance: u8,
-    remand: Mutex<HashMap<IpAddr, u8>>,
+    jailtime: u32,
+    remand: Mutex<HashMap<IpAddr, (u8, u64)>>,
+}
+
+fn exec(cmd: &str, err: &str) -> Result<(), Error> {
+    let sentence_sl: Vec<&str> = cmd.split_whitespace().collect();
+    let out = Command::new("sudo").args(sentence_sl).output()?;
+    let sc = out.status.code();
+    ensure!(sc == Some(0), "err exec {}, {:?}\n{}", cmd, out, err);
+    Ok(())
 }
 
 impl Jail {
     pub fn new(allowance: u8, jailtime: u32) -> Result<Jail> {
-        const ERR_MSG: &str = "error using ipset/iptables, maybe it's not installed, this program isn't running as root ?";
+        const ERR_MSG: &str = "error using ipset/iptables, maybe it's not installed, or this program isn't running as root ?";
         let n = format!("blockfast_jail_{}", jailtime);
 
-        let i0 = format!("ipset create -exist {} hash:ip timeout {}", n, jailtime);
-        let i1 = format!("iptables -I INPUT 1 -m set -j DROP --match-set {} src", n);
-        let i2 = format!("iptables -I FORWARD 1 -m set -j DROP --match-set {} src", n);
-
-        let args0: Vec<&str> = i0.split_whitespace().collect();
-        let args1: Vec<&str> = i1.split_whitespace().collect();
-        let args2: Vec<&str> = i2.split_whitespace().collect();
-
         // create
-        let out = Command::new("sudo").args(args0).output()?;
-        ensure!(out.status.code() == Some(0), "{}: {:?}", ERR_MSG, out);
+        let cmd = format!("ipset create -exist {} hash:ip timeout {}", n, jailtime);
+        exec(&cmd, ERR_MSG)?;
 
         // setup input
-        let out = Command::new("sudo").args(args1).output()?;
-        ensure!(out.status.code() == Some(0), "{}: {:?}", ERR_MSG, out);
+        let cmd = format!("iptables -I INPUT 1 -m set -j DROP --match-set {} src", n);
+        exec(&cmd, ERR_MSG)?;
 
         // setup fwd
-        let out = Command::new("sudo").args(args2).output()?;
-        ensure!(out.status.code() == Some(0), "{}: {:?}", ERR_MSG, out);
+        let cmd = format!("iptables -I FORWARD 1 -m set -j DROP --match-set {} src", n);
+        exec(&cmd, ERR_MSG)?;
 
         log!("jail setup, allowance {}, time {}s", allowance, jailtime);
         Ok(Jail {
             name: n,
             allowance,
+            jailtime,
             remand: Mutex::new(HashMap::new()),
         })
     }
 
     pub fn sentence(&self, ip: IpAddr, target: &str) -> Result<()> {
+        let now = get_epoch();
+
         let should_ban = {
             let mut locked_map = self.remand.lock().map_err(|_| anyhow!("cant lock"))?;
 
-            // TODO: set time of last offence, and add grace
-            let hits = *locked_map.entry(ip).and_modify(|e| *e += 1).or_insert(1);
+            let (hits, _ts) = *locked_map
+                .entry(ip)
+                .and_modify(|(hits, ts)| {
+                    if *ts + self.jailtime as u64 > now { // reset if we have a hit, but past the defined jailtime
+                        *ts = now;
+                        *hits = 1;
+                    } else {
+                        *hits += 1; // bump
+                    }
+                })
+                .or_insert((1, now));
 
             if hits < self.allowance {
                 false
@@ -63,12 +76,8 @@ impl Jail {
 
         if should_ban {
             log!("{} jailtime for: {}", target, ip);
-            let sentence = format!("ipset add -exist {} {}", self.name, ip);
-            let sentence_sl: Vec<&str> = sentence.split_whitespace().collect();
-
-            let out = Command::new("sudo").args(sentence_sl).output()?;
-            let stderr = std::str::from_utf8(&out.stderr)?;
-            ensure!(out.status.code() == Some(0), "executing ban {}", stderr);
+            let cmd = format!("ipset add -exist {} {}", self.name, ip);
+            exec(&cmd, "")?;
         }
 
         Ok(())

src/utils.rs

@@ -7,15 +7,20 @@ pub enum ParsingStatus {
     BadEntry(IpAddr),
 }
 
+pub fn get_epoch() -> u64 {
+    let e = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH);
+    e.map(|e| e.as_secs()).unwrap_or(0)
+}
+
 macro_rules! log{
     ($first:expr) => {
-        let e = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?;
-        eprintln!("{} ~ {}", e.as_secs(), $first);
+        let ts = crate::utils::get_epoch();
+        eprintln!("{} ~ {}", ts, $first);
     };
     ($first:expr, $($others:expr),+) => {
-        let e = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?;
+        let ts = crate::utils::get_epoch();
         let formatted = format!($first, $($others), *);
-        eprintln!("{} ~ {}", e.as_secs(), formatted);
+        eprintln!("{} ~ {}", ts, formatted);
     };
 }