diff --git a/src/jail.rs b/src/jail.rs index fe0f8dd..ba77f9d 100644 --- a/src/jail.rs +++ b/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>, + jailtime: u32, + remand: Mutex>, +} + +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 { - 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(()) diff --git a/src/utils.rs b/src/utils.rs index 61c82d5..44c78ea 100644 --- a/src/utils.rs +++ b/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); }; }