diff --git a/src/clf.rs b/src/clf.rs index 972e27f..63453ad 100644 --- a/src/clf.rs +++ b/src/clf.rs @@ -1,3 +1,4 @@ +use crate::utils::ParsingStatus; use anyhow::Result; use lazy_static::lazy_static; use std::net::IpAddr; @@ -8,7 +9,7 @@ lazy_static! { static ref BAD_STATUSES: [u32; 2] = [401, 429]; } -pub fn parse(line: &str) -> Result> { +pub fn parse(line: &str) -> Result { // TODO: Use a proper parser ? let elts: Vec<&str> = line.split_whitespace().collect(); @@ -20,11 +21,11 @@ pub fn parse(line: &str) -> Result> { for status in BAD_STATUSES.iter() { if *status == http_code { - return Ok(Some(ip)); + return Ok(ParsingStatus::BadEntry(ip)); } } - Ok(None) + Ok(ParsingStatus::OkEntry) } #[cfg(test)] @@ -39,8 +40,11 @@ mod tests { ]; vectors.iter().for_each(|e| { - let ret = parse(*e); - assert!(ret.unwrap().is_some()); + let ret = parse(*e).unwrap(); + match ret { + ParsingStatus::BadEntry(_) => {} + _ => panic!("bad parsing"), + } }) } @@ -52,8 +56,11 @@ mod tests { ]; vectors.iter().for_each(|e| { - let ret = parse(*e); - assert!(ret.unwrap().is_none()); + let ret = parse(*e).unwrap(); + match ret { + ParsingStatus::OkEntry => {} + _ => panic!("bad parsing"), + } }) } diff --git a/src/jail.rs b/src/jail.rs index f64e17c..1038c09 100644 --- a/src/jail.rs +++ b/src/jail.rs @@ -5,6 +5,8 @@ use std::sync::Mutex; use anyhow::*; +use crate::utils::JailStatus; + pub struct Jail { jailtime: u32, allowance: u8, @@ -91,7 +93,7 @@ impl Jail { }) } - pub fn probe(&self, ip: IpAddr) -> Result { + pub fn probe(&self, ip: IpAddr) -> Result { let should_ban = { let mut locked_map = self.remand.lock().map_err(|_| anyhow!("cant lock"))?; @@ -108,9 +110,9 @@ impl Jail { if should_ban { ipset_block(self.jailtime, ip)?; - Ok(true) + Ok(JailStatus::Jailed(ip)) } else { - Ok(false) + Ok(JailStatus::Remand) } } } diff --git a/src/main.rs b/src/main.rs index 037988b..92b1833 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,8 +7,15 @@ mod utils; mod jail; use crate::jail::Jail; +use crate::utils::*; -fn judge(path_sshd: &str, path_clf: &str, payload: &str, path: &str, jail: &Jail) -> Result<()> { +fn judge( + path_sshd: &str, + path_clf: &str, + payload: &str, + path: &str, + jail: &Jail, +) -> Result { let do_sshd = !path_sshd.is_empty(); let do_clf = !path_clf.is_empty(); let mut target = ""; @@ -24,16 +31,16 @@ fn judge(path_sshd: &str, path_clf: &str, payload: &str, path: &str, jail: &Jail }; let ip = match ret_parse? { - Some(ip) => ip, - None => return Ok(()), + ParsingStatus::OkEntry => return Ok(Judgment::Good), + ParsingStatus::BadEntry(ip) => ip, }; - if jail.probe(ip)? { - eprintln!("~ {} - too many infraction, jailtime for: {}", target, ip); + match jail.probe(ip)? { + JailStatus::Remand => Ok(Judgment::Remand), + JailStatus::Jailed(ip) => Ok(Judgment::Bad(target, ip)), } - - Ok(()) } + async fn run() -> Result<()> { let args = utils::cli().get_matches(); let mut lines = MuxedLines::new()?; @@ -69,9 +76,14 @@ async fn run() -> Result<()> { let payload = line.line(); let path = line.source().display().to_string(); - if let Err(err) = judge(path_sshd, path_clf, payload, &path, &jail) { - eprintln!("! ERR {:?} - file {}", err, path) - } + 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!("~ {} jailtime for: {}", target, ip) + } + }; } Ok(()) diff --git a/src/sshd.rs b/src/sshd.rs index 2d19bda..3667bd7 100644 --- a/src/sshd.rs +++ b/src/sshd.rs @@ -1,10 +1,11 @@ use anyhow::*; -use std::net::IpAddr; - use lazy_static::lazy_static; use regex::Regex; +use std::net::IpAddr; use std::str::FromStr; +use crate::utils::ParsingStatus; + struct Rule { matcher: String, extractor: Regex, @@ -27,7 +28,7 @@ lazy_static! { ]; } -pub fn parse(line: &str) -> Result> { +pub fn parse(line: &str) -> Result { let hits = SSHD_BAD.iter().find_map(|rule| { if line.contains(&rule.matcher) { rule.extractor.captures(line) @@ -37,16 +38,17 @@ pub fn parse(line: &str) -> Result> { }); if hits.is_none() { - return Ok(None); + return Ok(ParsingStatus::OkEntry); } let ip = hits .and_then(|c| c.get(2)) - .and_then(|m| IpAddr::from_str(m.as_str()).ok()); + .and_then(|m| Some(m.as_str())) + .ok_or(anyhow!("sshd cant extract ip"))?; - match ip { - Some(ip) => Ok(Some(ip)), - None => Err(anyhow!("cant parse sshd entry")), + match IpAddr::from_str(ip) { + Ok(ip) => Ok(ParsingStatus::BadEntry(ip)), + Err(_) => Err(anyhow!("cant parse sshd entry")), } } @@ -63,8 +65,11 @@ mod tests { ]; vectors.iter().for_each(|e| { - let ret = parse(*e); - assert!(ret.unwrap().is_some()); + let ret = parse(*e).unwrap(); + match ret { + ParsingStatus::BadEntry(_) => {} + _ => panic!("bad parsing"), + } }) } @@ -77,8 +82,11 @@ mod tests { ]; vectors.iter().for_each(|e| { - let ret = parse(*e); - assert!(ret.unwrap().is_none()); + let ret = parse(*e).unwrap(); + match ret { + ParsingStatus::OkEntry => {} + _ => panic!("bad parsing"), + } }) } diff --git a/src/utils.rs b/src/utils.rs index 596ab87..afd98a8 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,4 +1,20 @@ use clap::{App, Arg}; +use std::net::IpAddr; + +pub enum ParsingStatus { + OkEntry, + BadEntry(IpAddr), +} +pub enum Judgment { + Good, + Remand, + Bad(&'static str, IpAddr), +} + +pub enum JailStatus { + Remand, + Jailed(IpAddr), +} pub fn cli() -> App<'static, 'static> { App::new("ban internets scanner fast 🍶")