improving martillo-maldito to server to protect also nodes (port, ip, etc..)

src/iptables_wrapper.rs

@@ -1,4 +1,7 @@
+use std::collections::HashMap;
+
 use iptables::IPTables;
+use regex::Regex;
 
 pub fn is_port_secured(port: u16, docker: bool) -> bool {
     let iptables = iptables::new(false).unwrap();
@@ -15,6 +18,22 @@ pub fn is_port_secured(port: u16, docker: bool) -> bool {
     false
 }
 
+pub fn list_secured_ports(docker: bool) -> Vec<u16> {
+    let iptables = iptables::new(false).unwrap();
+    let rules = iptables.list("filter", &get_chain(docker));
+    if rules.is_err() {
+        return vec![];
+    }
+
+    let rgx = get_regex_for_port();
+    rules
+        .unwrap()
+        .iter()
+        .filter(|r| r.contains("-p tcp -m tcp --dport") && r.contains("-j DROP"))
+        .map(|r| extract_port(&rgx, r).unwrap())
+        .collect()
+}
+
 pub fn list_banned_ips() -> Vec<String> {
     let iptables = iptables::new(false).unwrap();
     let rules = iptables.list("filter", &get_chain(false));
@@ -22,66 +41,98 @@ pub fn list_banned_ips() -> Vec<String> {
         return vec![];
     }
 
-    // TODO: Remove after testing
-    println!("{:?}", rules);
-
+    let rgx = get_regex_for_ip();
     rules
         .unwrap()
         .iter()
-        .filter(|r| r.contains("--source") && r.contains("-j DROP"))
-        .map(|r| r.to_string())
+        .filter(|r| r.contains("-A INPUT") && r.contains("-j DROP") && r.contains("-s"))
+        .map(|r| extract_ip(&rgx, r).unwrap())
         .collect()
 }
 
-pub fn ban_port(port: u16, docker: bool, position: Option<i32>) {
+pub fn map_secured_ports_allowed_ips(docker: bool) -> HashMap<u16, Vec<String>> {
+    let mut result: HashMap<u16, Vec<String>> = HashMap::new();
+    let secured_ports = list_secured_ports(docker);
+    if secured_ports.is_empty() {
+        return result;
+    }
+
+    let iptables = iptables::new(false).unwrap();
+    let rules = iptables.list("filter", &get_chain(docker));
+    if rules.is_err() {
+        return result;
+    }
+
+    let rules = rules.unwrap();
+    let rgx = get_regex_for_ip();
+    for port in secured_ports {
+        let ips = rules
+            .iter()
+            .filter(|r| {
+                r.contains("-A INPUT -s")
+                    && r.contains(&format!("-p tcp -m tcp --dport {} -j ACCEPT", port))
+            })
+            .map(|r| extract_ip(&rgx, r).unwrap())
+            .collect();
+        result.insert(port, ips);
+    }
+
+    result
+}
+
+pub fn secure_port_rule(port: u16) -> String {
+    format!("-p tcp --dport {} -j DROP", port)
+}
+
+pub fn secure_port(port: u16, docker: bool, position: Option<i32>) {
     let iptables = iptables::new(false).unwrap();
 
     let table = "filter";
     let chain = get_chain(docker);
-    let rule = format!("-p tcp --dport {} -j DROP", port);
+    let rule = secure_port_rule(port);
     if let Some(position) = position {
         insert_unique(&iptables, table, &chain, &rule, position)
     } else {
         append_unique(&iptables, table, &chain, &rule)
     }
 
-    println!("banned port {}", port);
+    println!("Port {} secured", port);
 }
 
-pub fn unban_port(port: u16, docker: bool) {
+pub fn unsecure_port(port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
-    let _ = iptables.delete(
-        "filter",
-        &get_chain(docker),
-        &format!("-p tcp --dport {} -j DROP", port),
-    );
+    let _ = iptables.delete("filter", &get_chain(docker), &secure_port_rule(port));
 
-    println!("unbanned port {}", port);
+    println!("Port {} unsecured", port);
 }
 
-pub fn allow_ip_port(ip: &str, port: u16, docker: bool, position: Option<i32>) {
+pub fn allow_ip_for_port_rule(port: u16, ip: &str) -> String {
+    format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip)
+}
+
+pub fn allow_ip_for_port(ip: &str, port: u16, docker: bool, position: Option<i32>) {
     let iptables = iptables::new(false).unwrap();
 
     let table = "filter";
     let chain = get_chain(docker);
-    let rule = format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip);
+    let rule = allow_ip_for_port_rule(port, ip);
     if let Some(position) = position {
         insert_unique(&iptables, table, &chain, &rule, position)
     } else {
         append_unique(&iptables, table, &chain, &rule)
     }
 
-    println!("allowed {} to access {}", ip, port);
+    println!("Allowed {} to access {}", ip, port);
 }
 
-pub fn remove_ip_port(ip: &str, port: u16, docker: bool) {
+pub fn remove_allow_ip_for_port(ip: &str, port: u16, docker: bool) {
     let iptables = iptables::new(false).unwrap();
     let _ = iptables.delete(
         "filter",
         &get_chain(docker),
-        &format!("-p tcp --dport {} -s {} -j ACCEPT", port, ip),
+        &allow_ip_for_port_rule(port, ip),
     );
-    println!("removed access {} to {}", ip, port);
+    println!("Removed access of {} to {}", ip, port);
 }
 
 fn get_chain(docker: bool) -> String {
@@ -99,3 +150,42 @@ fn append_unique(iptables: &IPTables, table: &str, chain: &str, rule: &str) {
 fn insert_unique(iptables: &IPTables, table: &str, chain: &str, rule: &str, position: i32) {
     let _ = iptables.insert_unique(table, chain, rule, position);
 }
+
+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 get_regex_for_ip() -> Regex {
+    Regex::new(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b").unwrap()
+}
+
+fn get_regex_for_port() -> Regex {
+    Regex::new(r"--dport\s+(\d+)").unwrap()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn correct_extract_ip() {
+        let regex = get_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_extrac_ip() {
+        let regex = get_regex_for_ip();
+        let input = "-A INPUT -j DROP";
+        assert_eq!(extract_ip(&regex, input), None);
+    }
+}