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

use cli::Cli;
use iptables_wrapper::{
    allow_ip_for_port, is_port_secured, list_banned_ips, list_secured_ports,
    map_secured_ports_allowed_ips, remove_allow_ip_for_port, secure_port, unsecure_port,
};
use log::{error, info};
use login_attempt::LoginAttempt;
use std::{
    collections::HashMap,
    io::BufRead,
    process::{Child, Command, Stdio},
    thread::{sleep, spawn},
    time::Duration,
};
use structopt::StructOpt;

#[tokio::main]
async fn main() {
    start_logger();

    match Cli::from_args() {
        Cli::BanServer => {
            if let Err(error) = start_ban_server().await {
                error!("ban server: {error}");
            }
        }
        Cli::ListBannedIps { docker } => {
            println!(
                "{}",
                serde_json::to_string(&list_banned_ips(docker)).unwrap()
            );
        }
        Cli::ListSecuredPorts { docker } => {
            println!(
                "{}",
                serde_json::to_string(&list_secured_ports(docker)).unwrap()
            );
        }
        Cli::MapSecuredPortsAllowedIps { docker } => {
            println!(
                "{}",
                serde_json::to_string(&map_secured_ports_allowed_ips(docker)).unwrap()
            );
        }
        Cli::IsPortSecured { port, docker } => {
            println!("{}", is_port_secured(port, docker));
        }
        Cli::SecurePort {
            port,
            docker,
            position,
        } => println!("{}", secure_port(port, docker, position).is_ok()),
        Cli::UnsecurePort { port, docker } => println!("{}", unsecure_port(port, docker).is_ok()),
        Cli::AllowIpForPort {
            ip,
            port,
            docker,
            position,
        } => println!("{}", allow_ip_for_port(&ip, port, docker, position).is_ok()),
        Cli::OnlyIpForPort { ip, port, docker } => {
            let allowed = allow_ip_for_port(&ip, port, docker, Some(1));
            let secured = secure_port(port, docker, Some(2));
            println!("{}", allowed.is_ok() && secured.is_ok());
        }
        Cli::RemoveAllowIpPort { ip, port, docker } => {
            println!("{}", remove_allow_ip_for_port(&ip, port, docker).is_ok())
        }
        Cli::SaveIPTables => {
            println!("{}", iptables_save::save_iptables().is_ok())
        }
    }
}

async fn start_ban_server() -> std::io::Result<()> {
    let seconds_iptables = Duration::from_secs(60);
    info!("saving IPTables every {} secs", seconds_iptables.as_secs());

    spawn(move || loop {
        sleep(seconds_iptables);
        iptables_save::save_iptables();
    });

    let mut child: Child = Command::new("journalctl")
        .arg("-D")
        .arg("/var/log/journal")
        .arg("-u")
        .arg("ssh")
        .arg("-f")
        .stdout(Stdio::piped())
        .spawn()?;

    let stdout = child.stdout.as_mut().expect("Failed to capture stdout");
    let mut reader = std::io::BufReader::new(stdout);

    let iptables = iptables::new(false).unwrap();
    let mut login_attempts: HashMap<String, usize> = HashMap::new();

    info!("Start reading logins from SSH");
    loop {
        let mut line = String::new();

        if let Err(err) = reader.read_line(&mut line) {
            error!("Reading line: {}", err);
            continue;
        }

        println!("line: {}", line);

        if line.is_empty() {
            continue;
        }

        if let Some(login_attempt) = LoginAttempt::capture(&line) {
            info!(
                "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 {
                        if iptables
                            .append_unique(
                                "filter",
                                "INPUT",
                                &format!("--source {} -j DROP", login_attempt.ip),
                            )
                            .is_ok()
                        {
                            info!("IP {} banned", login_attempt.ip);
                        }

                        login_attempts.remove(&login_attempt.ip);
                    }
                }
                None => {
                    login_attempts.insert(login_attempt.ip, 1);
                }
            }
        }
    }
}

fn start_logger() {
    unsafe {
        std::env::set_var(
            "RUST_LOG",
            std::env::var("RUST_LOG").unwrap_or("INFO".to_string()),
        );
    }
    logfmt_logger::init();
}