cleanup & optimisation

src/jail.rs

@@ -5,95 +5,48 @@ use std::sync::Mutex;
 
 use anyhow::*;
 
-use crate::utils::JailStatus;
+use crate::utils::log;
 
 pub struct Jail {
-    jailtime: u32,
+    name: String,
     allowance: u8,
     remand: Mutex<HashMap<IpAddr, u8>>,
 }
 
-const JAIL_NAME: &str = "blockfast_jail";
-
-const ERR_MSG: &str =
-    "error using ipset/iptables, maybe it's not installed, this program isn't running as root ?";
-
-fn ipset_init() -> Result<()> {
-    let init0 = format!("ipset create {} hash:ip timeout 0", JAIL_NAME);
-    let init1 = format!(
-        "iptables -I INPUT 1 -m set -j DROP --match-set {} src",
-        JAIL_NAME
-    );
-    let init2 = format!(
-        "iptables -I FORWARD 1 -m set -j DROP --match-set {} src",
-        JAIL_NAME
-    );
-
-    let args0: Vec<&str> = init0.split_whitespace().collect();
-    let args1: Vec<&str> = init1.split_whitespace().collect();
-    let args2: Vec<&str> = init2.split_whitespace().collect();
-
-    // create
-    let out = Command::new("sudo").args(args0).output()?;
-    if out.status.code() != Some(0) {
-        let already_exists =
-            std::str::from_utf8(&out.stderr)?.contains("set with the same name already exists");
-
-        if already_exists {
-            return Ok(());
-        } else {
-            eprintln!("{:?}", out);
-            bail!(ERR_MSG);
-        }
-    }
-
-    // setup input
-    let out = Command::new("sudo").args(args1).output()?;
-    if out.status.code() != Some(0) {
-        eprintln!("{:?}", out);
-        bail!(ERR_MSG);
-    }
-
-    // setup fwd
-    let out = Command::new("sudo").args(args2).output()?;
-    if out.status.code() != Some(0) {
-        eprintln!("{:?}", out);
-        bail!(ERR_MSG);
-    }
-
-    Ok(())
-}
-
-fn ipset_block(jailtime: u32, ip: IpAddr) -> Result<()> {
-    let sentence = format!(
-        "ipset add {} {} timeout {}",
-        JAIL_NAME,
-        ip.to_string(),
-        jailtime
-    );
-    let sentence_sl: Vec<&str> = sentence.split_whitespace().collect();
-
-    let out = Command::new("sudo").args(sentence_sl).output()?;
-    if out.status.code() != Some(0) {
-        eprintln!("{:?}", out);
-        bail!("error executing ipset ban");
-    }
-
-    Ok(())
-}
-
 impl Jail {
     pub fn new(allowance: u8, jailtime: u32) -> Result<Jail> {
-        ipset_init()?;
+        const ERR_MSG: &str = "error using ipset/iptables, maybe it's not installed, 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);
+
+        // setup input
+        let out = Command::new("sudo").args(args1).output()?;
+        ensure!(out.status.code() == Some(0), "{}: {:?}", ERR_MSG, out);
+
+        // setup fwd
+        let out = Command::new("sudo").args(args2).output()?;
+        ensure!(out.status.code() == Some(0), "{}: {:?}", ERR_MSG, out);
+
+        log!("jail setup, allowance {}, time {}s", allowance, jailtime);
         Ok(Jail {
+            name: n,
             allowance,
-            jailtime,
             remand: Mutex::new(HashMap::new()),
         })
     }
 
-    pub fn probe(&self, ip: IpAddr) -> Result<JailStatus> {
+    pub fn sentence(&self, ip: IpAddr, target: &str) -> Result<()> {
         let should_ban = {
             let mut locked_map = self.remand.lock().map_err(|_| anyhow!("cant lock"))?;
 
@@ -109,10 +62,15 @@ impl Jail {
         };
 
         if should_ban {
-            ipset_block(self.jailtime, ip)?;
-            Ok(JailStatus::Jailed(ip))
-        } else {
-            Ok(JailStatus::Remand)
+            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);
         }
+
+        Ok(())
     }
 }