WIP moving code into lib

2024-12-31 03:57:40 +01:00
parent 41abf04a63
commit 20f1038ff8
9 changed files with 448 additions and 420 deletions
@@ -0,0 +1,243 @@
+use regex::Regex;
+use std::{collections::HashMap, process::Command};
+
+pub struct MartilloMaldito {
+    iptables: iptables::IPTables,
+    chain: String,
+}
+
+impl MartilloMaldito {
+    pub fn ipv4(docker: bool) -> MartilloMaldito {
+        MartilloMaldito {
+            iptables: iptables::new(false).unwrap(),
+            chain: Self::get_chain(docker).to_string(),
+        }
+    }
+
+    pub fn ipv6(docker: bool) -> MartilloMaldito {
+        MartilloMaldito {
+            iptables: iptables::new(true).unwrap(),
+            chain: Self::get_chain(docker).to_string(),
+        }
+    }
+
+    pub fn save_rules() -> std::io::Result<std::process::Output> {
+        Command::new("iptables-save")
+            .args(["-f", "/etc/iptables/rules.v4"])
+            .output()
+    }
+
+    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"))
+            .map(|r| extract_port(&port_regex, r).unwrap())
+            .collect()
+    }
+
+    pub fn get_banned_ips(&self) -> Vec<String> {
+        let rules = self.get_rules();
+        if rules.is_err() {
+            return vec![];
+        }
+
+        let ip_regex = iptables_regex_for_ip();
+        rules
+            .unwrap()
+            .iter()
+            .filter(|r| {
+                r.contains(&format!("-A {}", self.chain))
+                    && r.contains("-j DROP")
+                    && r.contains("-s")
+            })
+            .map(|r| extract_ip(&ip_regex, r).unwrap())
+            .collect()
+    }
+
+    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.append_unique("filter", &format!("-s {} -j DROP", ip))
+    }
+
+    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()
+}
+
+#[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()));
+    }
+}