1 files changed, 74 insertions, 58 deletions

diff --git a/src/fw.rs b/src/fw.rs
index 355e928..df3322e 100644
--- a/src/fw.rs
+++ b/src/fw.rs
@@ -1,6 +1,6 @@
 // use std::collections::HashSet;
 
-use anyhow::{Result,Context};
+use anyhow::{Context, Result};
 use iptables;
 use log::*;
 use regex::Regex;
@@ -8,77 +8,93 @@ use regex::Regex;
 use crate::messages;
 
 pub fn setup(ipt: &iptables::IPTables) -> Result<()> {
+  // ensure we start from a clean state without any rule already set
+  cleanup(ipt)?;
 
-    // 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")?;
+  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(())
+  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(())
+  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::new();
-    // 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")
-                }
-
-            },
-            _ => {}
+  let mut ports = messages::PublicExposedPorts::new();
+  // 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)
+  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")?;
+  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")?;
     }
 
-    Ok(())
-}
+    ipt
+      .delete_chain("filter", "DIPLONAT")
+      .context("Failed to delete chain")?;
+  }
 
+  Ok(())
+}