martillo-maldito/src/martillo_maldito.rs

use crate::ipset::{Ipset, DEFAULT_SET_NAME, DEFAULT_TIMEOUT_SECS};
use regex::Regex;
use std::{collections::HashMap, process::Command};

pub struct MartilloMaldito {
    iptables: iptables::IPTables,
    chain: String,
    ipset: Ipset,
    ban_timeout: Option<u32>,
}

impl MartilloMaldito {
    pub fn ipv4(docker: bool) -> MartilloMaldito {
        Self::ipv4_with_config(docker, DEFAULT_SET_NAME, None)
    }

    pub fn ipv4_with_config(
        docker: bool,
        ipset_name: &str,
        ban_timeout: Option<u32>,
    ) -> MartilloMaldito {
        let ipset = Ipset::new(ipset_name);
        if let Err(e) = ipset.ensure_exists() {
            eprintln!("Warning: could not ensure ipset exists: {}", e);
        }
        let martillo = MartilloMaldito {
            iptables: iptables::new(false).unwrap(),
            chain: Self::get_chain(docker).to_string(),
            ipset,
            ban_timeout,
        };
        if let Err(e) = martillo.ensure_ipset_match_rule() {
            eprintln!("Warning: could not install ipset match rule: {}", e);
        }
        martillo
    }

    fn ensure_ipset_match_rule(&self) -> Result<(), Box<dyn std::error::Error>> {
        let match_rule = format!("-m set --match-set {} src -j DROP", self.ipset.name());
        let args: Vec<&str> = match_rule.split_whitespace().collect();
        let check_args: Vec<&str> = std::iter::once("-C").chain(self.chain.split_whitespace()).chain(args.iter().copied()).collect();
        let output = Command::new("iptables").args(&check_args).output()?;
        if output.status.success() {
            return Ok(());
        }
        let insert_args: Vec<&str> = std::iter::once("-I").chain(self.chain.split_whitespace()).chain(std::iter::once("1")).chain(args.iter().copied()).collect();
        let output = Command::new("iptables").args(&insert_args).output()?;
        if !output.status.success() {
            return Err(format!(
                "failed to install ipset match rule: {}",
                String::from_utf8_lossy(&output.stderr)
            )
            .into());
        }
        Ok(())
    }

    pub fn save_rules() -> std::io::Result<std::process::Output> {
        Command::new("iptables-save").args(["-f", "/etc/iptables/rules.v4"]).output()
    }

    pub fn save_ipset(&self, path: &std::path::Path) -> Result<(), Box<dyn std::error::Error>> {
        self.ipset.save(path)?;
        Ok(())
    }

    pub fn is_port_secured(&self, port: u16) -> bool {
        let rules = self.get_rules();
        if rules.is_err() {
            return false;
        }
        for rule in rules.unwrap() {
            if rule.contains(&format!("-p tcp -m tcp --dport {} -j DROP", port)) {
                return true;
            }
        }
        false
    }

    pub fn get_secured_ports(&self) -> Vec<u16> {
        let rules = self.get_rules();
        if rules.is_err() {
            return vec![];
        }
        let port_regex = iptables_regex_for_port();
        rules
            .unwrap()
            .iter()
            .filter(|r| r.contains("-p tcp -m tcp --dport") && r.contains("-j DROP"))
            .filter(|r| !r.contains("match-set"))
            .map(|r| extract_port(&port_regex, r).unwrap())
            .collect()
    }

    pub fn get_banned_ips(&self) -> Vec<String> {
        self.ipset.list().unwrap_or_default()
    }

    pub fn get_secured_ports_with_allowed_ips(&self) -> HashMap<u16, Vec<String>> {
        let mut result: HashMap<u16, Vec<String>> = HashMap::new();
        let secured_ports = self.get_secured_ports();
        if secured_ports.is_empty() {
            return result;
        }
        let rules = self.get_rules();
        if rules.is_err() {
            return result;
        }
        let rules = rules.unwrap();
        let ip_regex = iptables_regex_for_ip();
        for port in secured_ports {
            let ips = rules
                .iter()
                .filter(|r| {
                    r.contains(&format!("-A {} -s", self.chain))
                        && r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port))
                })
                .map(|r| extract_ip(&ip_regex, r).unwrap())
                .collect();
            result.insert(port, ips);
        }
        result
    }

    pub fn ban_ip(&self, ip: &str) -> Result<(), Box<dyn std::error::Error>> {
        self.ipset.add(ip, self.ban_timeout)?;
        Ok(())
    }

    pub fn ban_ip_with_timeout(
        &self,
        ip: &str,
        timeout_secs: u32,
    ) -> Result<(), Box<dyn std::error::Error>> {
        self.ipset.add(ip, Some(timeout_secs))?;
        Ok(())
    }

    pub fn unban_ip(&self, ip: &str) -> Result<(), Box<dyn std::error::Error>> {
        self.ipset.del(ip)?;
        Ok(())
    }

    pub fn secure_port(
        &self,
        port: u16,
        position: Option<usize>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let table = "filter";
        let rule = secure_port_rule(port);

        let position = if self.chain == "DOCKER-USER" && position.is_none() {
            let all_docker_rules = self.get_rules()?;
            Some(all_docker_rules.len() - 1)
        } else {
            position
        };

        if let Some(position) = position {
            self.insert_unique(table, &rule, position)
        } else {
            self.append_unique(table, &rule)
        }
    }

    pub fn unsecure_port(&self, port: u16) -> Result<(), Box<dyn std::error::Error>> {
        let rule = secure_port_rule(port);
        self.iptables.delete("filter", &self.chain, &rule)
    }

    pub fn allow_ip_for_port(
        &self,
        ip: &str,
        port: u16,
        position: Option<usize>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let table = "filter";
        let rule = allow_ip_for_port_rule(port, ip);

        if let Some(position) = position {
            self.insert_unique(table, &rule, position)
        } else {
            self.append_unique(table, &rule)
        }
    }

    pub fn remove_allow_ip_for_port(
        &self,
        ip: &str,
        port: u16,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let rule = allow_ip_for_port_rule(port, ip);
        self.iptables.delete("filter", &self.chain, &rule)
    }

    fn get_rules(&self) -> Result<Vec<String>, Box<dyn std::error::Error>> {
        self.iptables.list("filter", &self.chain)
    }

    fn append_unique(&self, table: &str, rule: &str) -> Result<(), Box<dyn std::error::Error>> {
        self.iptables.append_unique(table, &self.chain, rule)
    }

    fn insert_unique(
        &self,
        table: &str,
        rule: &str,
        position: usize,
    ) -> Result<(), Box<dyn std::error::Error>> {
        self.iptables
            .insert_unique(table, &self.chain, rule, position as i32)
    }

    fn get_chain(docker: bool) -> &'static str {
        if docker {
            "DOCKER-USER"
        } else {
            "INPUT"
        }
    }
}

fn secure_port_rule(port: u16) -> String {
    format!("-p tcp --dport {} -j DROP", port)
}

fn allow_ip_for_port_rule(port: u16, ip: &str) -> String {
    format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip)
}

fn extract_ip(regex: &Regex, input: &str) -> Option<String> {
    regex
        .captures(input)
        .and_then(|caps| caps.get(0).map(|m| m.as_str().to_string()))
}

fn extract_port(regex: &Regex, input: &str) -> Option<u16> {
    regex
        .captures(input)
        .and_then(|caps| caps.get(1).map(|m| m.as_str().parse::<u16>().unwrap()))
}

fn iptables_regex_for_ip() -> Regex {
    Regex::new(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b").unwrap()
}

fn iptables_regex_for_port() -> Regex {
    Regex::new(r"--dport\s+(\d+)").unwrap()
}

pub const fn default_ban_timeout_secs() -> u32 {
    DEFAULT_TIMEOUT_SECS
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn correct_extract_ip() {
        let regex = iptables_regex_for_ip();
        let input = "-A INPUT -s 81.69.255.132/32 -j DROP";
        assert_eq!(extract_ip(&regex, input), Some("81.69.255.132".to_string()));
    }

    #[test]
    fn no_match_extract_ip() {
        let regex = iptables_regex_for_ip();
        let input = "-A INPUT -j DROP";
        assert_eq!(extract_ip(&regex, input), None);
    }

    #[test]
    fn docker_extract_ip() {
        let regex = iptables_regex_for_ip();
        let input = "-A DOCKER -d 172.18.0.2/32 ! -i br-127d33df48a4 -o br-127d33df48a4 -p tcp -m tcp --dport 8078 -j ACCEPT";
        assert_eq!(extract_ip(&regex, input), Some("172.18.0.2".to_string()));
    }

    #[test]
    fn secured_ports_excludes_ipset_match_rule() {
        let input = "-A INPUT -p tcp -m tcp --dport 2222 -j DROP";
        assert_eq!(
            extract_port(&iptables_regex_for_port(), input),
            Some(2222)
        );
        assert!(!input.contains("match-set"));
    }

    #[test]
    fn default_timeout_constant_is_one_hour() {
        assert_eq!(default_ban_timeout_secs(), 3600);
    }
}