cleanup & optimisation

src/clf.rs

@@ -1,28 +1,35 @@
 use crate::utils::ParsingStatus;
-use anyhow::Result;
+use anyhow::*;
 use lazy_static::lazy_static;
-use std::net::IpAddr;
+use regex::Regex;
+use std::{net::IpAddr, str::FromStr};
 
-// TODO: allow user-provided list
-// TODO: match different error-levels (10 404, but only 5 401, etc...)
 lazy_static! {
     static ref BAD_STATUSES: [u32; 2] = [401, 429];
+    static ref RE_IP: Regex = Regex::new(r"^(\S+)\s").unwrap();
+    static ref RE_STATUS: Regex = Regex::new(r"(\d+)\s(\w+)$").unwrap();
 }
 
+#[allow(clippy::bind_instead_of_map)]
 pub fn parse(line: &str) -> Result<ParsingStatus> {
-    // TODO: Use a proper parser ?
-    let elts: Vec<&str> = line.split_whitespace().collect();
+    let ip = RE_IP
+        .captures(line)
+        .and_then(|c| c.get(1))
+        .and_then(|g| Some(g.as_str()))
+        .and_then(|e| IpAddr::from_str(e).ok())
+        .ok_or_else(|| anyhow!("cant parse clf line - ip"))?;
 
-    let ip_str = elts[0];
-    let ip = ip_str.parse::<IpAddr>()?;
+    let status = RE_STATUS
+        .captures(line)
+        .and_then(|c| c.get(1))
+        .and_then(|g| Some(g.as_str()))
+        .and_then(|e| e.parse::<u32>().ok())
+        .ok_or_else(|| anyhow!("cant parse clf line - status"))?;
 
-    let http_code_str = elts[elts.len() - 2] as &str;
-    let http_code = http_code_str.parse::<u32>()?;
+    let is_bad_status = BAD_STATUSES.iter().any(|s| s == &status);
 
-    for status in BAD_STATUSES.iter() {
-        if *status == http_code {
-            return Ok(ParsingStatus::BadEntry(ip));
-        }
+    if is_bad_status {
+        return Ok(ParsingStatus::BadEntry(ip));
     }
 
     Ok(ParsingStatus::OkEntry)

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(())
     }
 }

src/main.rs

@@ -1,5 +1,8 @@
+use std::path::PathBuf;
+use std::result::Result::Ok;
+
 use anyhow::*;
-use linemux::MuxedLines;
+use linemux::{Line, MuxedLines};
 
 mod clf;
 mod sshd;
@@ -9,41 +12,9 @@ mod jail;
 use crate::jail::Jail;
 use crate::utils::*;
 
-fn judge(
-    path_sshd: &str,
-    path_clf: &str,
-    payload: &str,
-    path: &str,
-    jail: &Jail,
-) -> Result<Judgment> {
-    let do_sshd = !path_sshd.is_empty();
-    let do_clf = !path_clf.is_empty();
-    let mut target = "";
-
-    let ret_parse = if do_sshd && path.ends_with(path_sshd) {
-        target = "sshd";
-        sshd::parse(payload)
-    } else if do_clf && path.ends_with(path_clf) {
-        target = "clf ";
-        clf::parse(payload)
-    } else {
-        Err(anyhow!("cant locate file !"))
-    };
-
-    let ip = match ret_parse? {
-        ParsingStatus::OkEntry => return Ok(Judgment::Good),
-        ParsingStatus::BadEntry(ip) => ip,
-    };
-
-    match jail.probe(ip)? {
-        JailStatus::Remand => Ok(Judgment::Remand),
-        JailStatus::Jailed(ip) => Ok(Judgment::Bad(target, ip)),
-    }
-}
-
 async fn run() -> Result<()> {
     let args = utils::cli().get_matches();
-    let mut lines = MuxedLines::new()?;
+    let mut ml = MuxedLines::new()?;
 
     // jail
     let jailtime_str = args.value_of("jailtime").unwrap_or("");
@@ -53,46 +24,54 @@ async fn run() -> Result<()> {
     let allowance = allowance_str.parse().context("parsing allowance")?;
 
     let jail = Jail::new(allowance, jailtime)?;
-    eprintln!(
-        "+ jail setup, offences allowed: {}, jailtime {}s",
-        allowance, jailtime
-    );
 
     // sshd
-    let path_sshd = args.value_of("sshd_logpath").unwrap_or("");
-    if !path_sshd.is_empty() {
-        lines.add_file(path_sshd).await?;
-        eprintln!("+ starting with sshd parsing at {}", path_sshd);
+    let mut path_sshd: PathBuf = args.value_of("sshd_logpath").unwrap_or("").into();
+    if path_sshd.exists() {
+        path_sshd = std::fs::canonicalize(path_sshd)?;
+        ml.add_file(&path_sshd).await?;
+        log!("starting with sshd parsing at {:?}", &path_sshd);
     }
 
     // common log format
-    let path_clf = args.value_of("clf_logpath").unwrap_or("");
-    if !path_clf.is_empty() {
-        lines.add_file(path_clf).await?;
-        eprintln!("+ starting with clf parsing at {}", path_clf);
+    let mut path_clf: PathBuf = args.value_of("clf_logpath").unwrap_or("").into();
+    if path_clf.exists() {
+        path_clf = std::fs::canonicalize(path_clf)?;
+        ml.add_file(&path_clf).await?;
+        log!("starting with clf parsing at {:?}", &path_clf);
     }
 
-    while let Ok(Some(line)) = lines.next_line().await {
+    let assess_line = |line: Line| {
         let payload = line.line();
-        let path = line.source().display().to_string();
+        let path = line.source();
 
-        match judge(path_sshd, path_clf, payload, &path, &jail) {
-            Err(err) => eprintln!("! ERR {:?} - file {}", err, path),
-            Ok(Judgment::Good) => {}
-            Ok(Judgment::Remand) => {}
-            Ok(Judgment::Bad(target, ip)) => {
-                eprintln!("~ too many infraction, {} jailtime for: {}", target, ip)
-            }
+        let (target, ret) = if path == path_sshd {
+            ("sshd", sshd::parse(payload)?)
+        } else if path == path_clf {
+            ("clf", clf::parse(payload)?)
+        } else {
+            bail!("file {:?} unknown", path)
         };
+
+        if let ParsingStatus::BadEntry(ip) = ret {
+            jail.sentence(ip, target)?;
+        }
+
+        Ok(())
+    };
+
+    while let Ok(Some(line)) = ml.next_line().await {
+        if let Err(e) = assess_line(line) {
+            log!("ERR: {:?}", e);
+        }
     }
 
     Ok(())
 }
 
 #[tokio::main]
-async fn main() -> std::io::Result<()> {
-    let ret = run().await;
-    let _ = ret.map_err(|e| eprintln!("! ERROR {:?}", e));
+async fn main() -> Result<()> {
+    run().await?;
     eprintln!("\n");
     let _ = utils::cli().print_help();
     Ok(())

src/sshd.rs

@@ -1,8 +1,7 @@
 use anyhow::*;
 use lazy_static::lazy_static;
 use regex::Regex;
-use std::net::IpAddr;
-use std::str::FromStr;
+use std::{net::IpAddr, str::FromStr};
 
 use crate::utils::ParsingStatus;
 
@@ -15,27 +14,24 @@ lazy_static! {
     static ref SSHD_BAD: [Rule; 3] = [
         Rule {
             matcher: "Failed password".to_string(),
-            extractor: Regex::new(r"(from.)(.*)(.port)").unwrap(),
+            extractor: Regex::new(r"(from.)(\S+)").unwrap(),
         },
         Rule {
             matcher: "Invalid user ".to_string(),
-            extractor: Regex::new(r"(from.)(.*)").unwrap(),
+            extractor: Regex::new(r"(from.)(\S+)").unwrap(),
         },
         Rule {
             matcher: "authentication failure".to_string(),
-            extractor: Regex::new(r"(rhost=)(.*)").unwrap()
+            extractor: Regex::new(r"(rhost=)(\S+)").unwrap()
         },
     ];
 }
 
 pub fn parse(line: &str) -> Result<ParsingStatus> {
-    let hits = SSHD_BAD.iter().find_map(|rule| {
-        if line.contains(&rule.matcher) {
-            rule.extractor.captures(line)
-        } else {
-            None
-        }
-    });
+    let hits = SSHD_BAD
+        .iter()
+        .find(|rule| line.contains(&rule.matcher))
+        .and_then(|r| r.extractor.captures(line));
 
     if hits.is_none() {
         return Ok(ParsingStatus::OkEntry);
@@ -43,12 +39,10 @@ pub fn parse(line: &str) -> Result<ParsingStatus> {
 
     let ip = hits
         .and_then(|c| c.get(2))
-        .and_then(|m| IpAddr::from_str(m.as_str()).ok());
+        .and_then(|m| IpAddr::from_str(m.as_str()).ok())
+        .ok_or_else(|| anyhow!("cant parse sshd line"))?;
 
-    match ip {
-        Some(ip) => Ok(ParsingStatus::BadEntry(ip)),
-        None => Err(anyhow!("cant parse sshd entry")),
-    }
+    Ok(ParsingStatus::BadEntry(ip))
 }
 
 #[cfg(test)]
@@ -93,12 +87,10 @@ mod tests {
     fn malformed() {
         let vectors = [
             "Sep 26 06:25:19 livecompute sshd[23246]: Failed password for root from 179.124.36.195.232 port 41883 ssh2",
-            "Sep 26 06:26:14 livecompute sshd[23292]: pam_unix(sshd:auth): authentication failure; logname= u =0 tty=ssh ruser= rhost=",
         ];
 
         vectors.iter().for_each(|e| {
-            let ret = parse(*e);
-            assert!(ret.is_err());
+            parse(*e).expect_err("");
         })
     }
 }

src/utils.rs

@@ -1,24 +1,29 @@
 use clap::{App, Arg};
 use std::net::IpAddr;
 
+#[derive(Debug)]
 pub enum ParsingStatus {
     OkEntry,
     BadEntry(IpAddr),
 }
-pub enum Judgment {
-    Good,
-    Remand,
-    Bad(&'static str, IpAddr),
+
+macro_rules! log{
+    ($first:expr) => {
+        let e = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?;
+        eprintln!("{} ~ {}", e.as_secs(), $first);
+    };
+    ($first:expr, $($others:expr),+) => {
+        let e = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?;
+        let formatted = format!($first, $($others), *);
+        eprintln!("{} ~ {}", e.as_secs(), formatted);
+    };
 }
 
-pub enum JailStatus {
-    Remand,
-    Jailed(IpAddr),
-}
+pub(crate) use log;
 
 pub fn cli() -> App<'static, 'static> {
     App::new("ban internets scanner fast 🍶")
-        .version("v0.0.1")
+        .version(env!("CARGO_PKG_VERSION"))
         .author("pierre dubouilh <pldubouilh@gmail.com>")
         // .arg(Arg::with_name("prune")
         //     .short("prune")
@@ -29,7 +34,7 @@ pub fn cli() -> App<'static, 'static> {
             Arg::with_name("jailtime")
                 .short("j")
                 .help("jail time (seconds)")
-                .default_value("3600")
+                .default_value("21600") // 6 hours
                 .takes_value(true),
         )
         .arg(
@@ -55,7 +60,7 @@ pub fn cli() -> App<'static, 'static> {
         )
     // .arg(Arg::with_name("clf_bad_http_codes")
     //     .short("cb")
-    //     .help("bad CLF http codes")
-    //     .default_value("{401, 429}")
+    //     .help("bad http statuses for CLF")
+    //     .default_value([401, 429])
     //     .takes_value(true))
 }