pub mod login_attempt;

use linemux::MuxedLines;
use login_attempt::LoginAttempt;
use regex::Regex;
use std::collections::HashMap;

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let iptables = iptables::new(false).unwrap();
    let mut lines = MuxedLines::new()?;
    lines.add_file("/host_ssh/auth.log").await?;
    let mut login_attempts: HashMap<String, usize> = HashMap::new();

    println!("Listening to changes over /host_ssh/auth.log");
    while let Ok(Some(line)) = lines.next_line().await {
        if let Some(login_attempt) = capture_failed_login_attempt(line.line()) {
            match login_attempts.get_mut(&login_attempt.ip) {
                Some(count) => {
                    *count += 1;

                    if *count == 3 {
                        iptables
                            .append_unique(
                                "filter",
                                "INPUT",
                                &format!("--source {} -j DROP", login_attempt.ip),
                            )
                            .unwrap();
                        login_attempts.remove(&login_attempt.ip);
                        println!("{} banned", login_attempt.ip);
                    }
                }
                None => {
                    login_attempts.insert(login_attempt.ip.clone(), 1);
                }
            }
        }
    }
    Ok(())
}

fn capture_failed_login_attempt(line: &str) -> Option<LoginAttempt> {
    let regex = Regex::new(r#"Failed password for (?:invalid user )?(?P<user>\S+) from (?P<ip>\S+) port (?P<port>\d+)"#).unwrap();
    let captured = regex.captures(line)?;
    Some(LoginAttempt::new(
        captured.name("ip").unwrap().as_str(),
        captured.name("user").unwrap().as_str(),
        captured.name("port").unwrap().as_str(),
    ))
}