pub mod cli;
pub mod iptables_save;
pub mod iptables_wrapper;
pub mod login_attempt;

use cli::Arguments;
use iptables_wrapper::{
    allow_ip_port, ban_port, is_port_secured, list_banned_ips, remove_ip_port, unban_port,
};
use linemux::MuxedLines;
use login_attempt::LoginAttempt;
use std::path::PathBuf;
use std::thread;
use std::{collections::HashMap, thread::sleep, time::Duration};
use structopt::StructOpt;

#[tokio::main]
async fn main() {
    let opts = Arguments::from_args();
    match opts {
        Arguments::BanServer {
            ssh_auth_log,
            iptables_save,
        } => {
            let _ = start_ban_server(ssh_auth_log, iptables_save).await;
        }
        Arguments::ListBannedIps => {
            for ip in list_banned_ips() {
                println!("{}", ip);
            }
        }
        Arguments::IsPortSecured { port, docker } => {
            let is_secured = is_port_secured(port, docker);
            println!("{}", is_secured);
        }
        Arguments::BanPort {
            port,
            docker,
            position,
        } => ban_port(port, docker, position),
        Arguments::UnbanPort { port, docker } => unban_port(port, docker),
        Arguments::AllowIpPort {
            ip,
            port,
            docker,
            position,
        } => allow_ip_port(&ip, port, docker, position),
        Arguments::OnlyIpPort { ip, port, docker } => {
            allow_ip_port(&ip, port, docker, Some(1));
            ban_port(port, docker, Some(2));
        }
        Arguments::RemoveIpPort { ip, port, docker } => remove_ip_port(&ip, port, docker),
    }
}

async fn start_ban_server(
    ssh_auth_log: PathBuf,
    iptables_save: Option<PathBuf>,
) -> std::io::Result<()> {
    let iptables = iptables::new(false).unwrap();
    let mut lines = MuxedLines::new()?;
    lines.add_file(&ssh_auth_log).await?;
    let mut login_attempts: HashMap<String, usize> = HashMap::new();

    if let Some(iptables_save) = iptables_save {
        let seconds_iptables = Duration::from_secs(60);
        println!(
            "starting iptables-save, save every {} seconds",
            seconds_iptables.as_secs()
        );
        thread::spawn(move || loop {
            sleep(seconds_iptables);
            iptables_save::save_iptables(&iptables_save);
        });
    }

    println!("listening to changes over {}", ssh_auth_log.display());
    while let Ok(Some(line)) = lines.next_line().await {
        if let Some(login_attempt) = LoginAttempt::capture(line.line()) {
            println!(
                "failed login attempt from {}@{}:{}",
                login_attempt.user, login_attempt.ip, login_attempt.port
            );
            match login_attempts.get_mut(&login_attempt.ip) {
                Some(count) => {
                    *count += 1;

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