1 files changed, 81 insertions, 0 deletions
diff --git a/src/fw.rs b/src/fw.rs
new file mode 100644
index 0000000..bc4d119
--- /dev/null
+++ b/src/fw.rs
@@ -0,0 +1,81 @@
+use iptables;
+use regex::Regex;
+use std::collections::HashSet;
+use crate::messages;
+use anyhow::{Result,Context};
+use log::*;
+
+pub fn setup(ipt: &iptables::IPTables) -> Result<()> {
+
+    // ensure we start from a clean state without any rule already set
+    cleanup(ipt)?;
+    
+    ipt.new_chain("filter", "DIPLONAT").context("Failed to create new chain")?;
+    ipt.insert_unique("filter", "INPUT", "-j DIPLONAT", 1).context("Failed to insert jump rule")?;
+
+    Ok(())
+}
+
+pub fn open_ports(ipt: &iptables::IPTables, ports: messages::PublicExposedPorts) -> Result<()> {
+    for p in ports.tcp_ports {
+        ipt.append("filter", "DIPLONAT", &format!("-p tcp --dport {} -j ACCEPT", p)).context("Failed to insert port rule")?;
+    }
+
+    for p in ports.udp_ports {
+        ipt.append("filter", "DIPLONAT", &format!("-p udp --dport {} -j ACCEPT", p)).context("Failed to insert port rule")?;
+    }
+
+    Ok(())
+}
+
+pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExposedPorts> {
+    let mut ports = messages::PublicExposedPorts {
+        tcp_ports: HashSet::new(),
+        udp_ports: HashSet::new()
+    };
+
+    let list = ipt.list("filter", "DIPLONAT")?;
+    let re = Regex::new(r"\-A.*? \-p (\w+).*\-\-dport (\d+).*?\-j ACCEPT").context("Regex matching open ports encountered an unexpected rule")?;
+    for i in list {
+        let caps = re.captures(&i);
+        match caps {
+            Some(c) => {
+           
+                if let (Some(raw_proto), Some(raw_port)) = (c.get(1), c.get(2)) {
+
+                    let proto = String::from(raw_proto.as_str());
+                    let number = String::from(raw_port.as_str()).parse::<u16>()?;
+
+                    if proto == "tcp" {
+                        ports.tcp_ports.insert(number);
+                    } else {
+                        ports.udp_ports.insert(number);
+                    }
+
+                } else {
+                    error!("Unexpected rule found in DIPLONAT chain")
+                }
+
+            },
+            _ => {}
+        }
+    }
+
+    Ok(ports)
+}
+
+pub fn cleanup(ipt: &iptables::IPTables) -> Result<()> {
+    
+    if ipt.chain_exists("filter", "DIPLONAT")? {
+        ipt.flush_chain("filter", "DIPLONAT").context("Failed to flush the DIPLONAT chain")?;
+    
+        if ipt.exists("filter", "INPUT", "-j DIPLONAT")? {
+            ipt.delete("filter", "INPUT", "-j DIPLONAT").context("Failed to delete jump rule")?;
+        }
+    
+        ipt.delete_chain("filter", "DIPLONAT").context("Failed to delete chain")?;
+    }
+
+    Ok(())
+}
+