1 files changed, 36 insertions, 31 deletions

diff --git a/src/fw.rs b/src/fw.rs
index 8ee3e6b..bc4d119 100644
--- a/src/fw.rs
+++ b/src/fw.rs
@@ -2,61 +2,58 @@ use iptables;
 use regex::Regex;
 use std::collections::HashSet;
 use crate::messages;
+use anyhow::{Result,Context};
+use log::*;
 
-#[derive(Debug)]
-pub struct FirewallError(String);
+pub fn setup(ipt: &iptables::IPTables) -> Result<()> {
 
-impl From<iptables::error::IPTError> for FirewallError {
-    fn from(error: iptables::error::IPTError) -> Self {
-        FirewallError(error.to_string())
-    }
-}
-
-pub fn setup(ipt: &iptables::IPTables) -> Result<(), FirewallError> {
-
-    if !ipt.chain_exists("filter", "DIPLONAT")? {
-        ipt.new_chain("filter", "DIPLONAT")?;
-    }
-
-    ipt.insert_unique("filter", "INPUT", "-j DIPLONAT", 1)?;
+    // 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<(), FirewallError> {
+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))?;
+        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))?;
+        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, FirewallError> {
+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").unwrap();
+    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) => {
-                let raw_proto = c.get(1).unwrap();
-                let raw_port = c.get(2).unwrap();
+           
+                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>()?;
 
-                let proto = String::from(raw_proto.as_str());
-                let number = String::from(raw_port.as_str()).parse::<u16>().unwrap();
+                    if proto == "tcp" {
+                        ports.tcp_ports.insert(number);
+                    } else {
+                        ports.udp_ports.insert(number);
+                    }
 
-                if proto == "tcp" {
-                    ports.tcp_ports.insert(number);
                 } else {
-                    ports.udp_ports.insert(number);
+                    error!("Unexpected rule found in DIPLONAT chain")
                 }
 
             },
@@ -67,10 +64,18 @@ pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExpo
     Ok(ports)
 }
 
-pub fn cleanup(ipt: &iptables::IPTables) -> Result<(), FirewallError> {
-    ipt.flush_chain("filter", "DIPLONAT")?;
-    ipt.delete("filter", "INPUT", "-j DIPLONAT")?;
-    ipt.delete_chain("filter", "DIPLONAT")?;
+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(())
 }