martillo-maldito/src/main.rs

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 linemux::MuxedLines;
use log::{error, info};
use login_attempt::LoginAttempt;
use std::{collections::HashMap, path::PathBuf, thread::sleep, thread::spawn, time::Duration};
use structopt::StructOpt;

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

    match Cli::from_args() {
        Cli::BanServer {
            ssh_auth_log,
            iptables_save,
        } => {
            let _ = start_ban_server(ssh_auth_log, iptables_save).await;
        }
        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 { iptables_save } => {
            let path = if let Some(iptables_save) = iptables_save {
                iptables_save
            } else {
                PathBuf::from("/etc/iptables/rules.v4")
            };

            println!("{}", iptables_save::save_iptables(&path).is_ok())
        }
    }
}

async fn start_ban_server(
    ssh_auth_log: PathBuf,
    iptables_save: Option<PathBuf>,
) -> std::io::Result<()> {
    if let Some(iptables_save) = iptables_save {
        let seconds_iptables = Duration::from_secs(60);
        println!(
            "Saving IPTables every {} seconds",
            seconds_iptables.as_secs()
        );

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

    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();

    info!("listening changes over {}", ssh_auth_log.display());
    loop {
        let next_line = lines.next_line().await;
        if let Err(err) = next_line {
            error!("reading next file: {}", err);
            continue;
        }

        let line = next_line?;
        if line.is_none() {
            continue;
        }

        let line = line.unwrap();
        if let Some(login_attempt) = LoginAttempt::capture(line.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);
                }
            }
        }
    }
}