1 files changed, 22 insertions, 15 deletions

diff --git a/src/fw.rs b/src/fw.rs
index 7650b3a..955425a 100644
--- a/src/fw.rs
+++ b/src/fw.rs
@@ -2,13 +2,7 @@ use iptables;
 use regex::Regex;
 use std::collections::HashSet;
 use std::io;
-
-
-#[derive(PartialEq,Eq,Debug,Hash)]
-pub struct Port {
-    proto: String,
-    number: u16,
-}
+use crate::messages;
 
 #[derive(Debug)]
 pub struct FirewallError(String);
@@ -17,26 +11,34 @@ 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> {
+    
     ipt.new_chain("filter", "DIPLONAT")?;
     ipt.insert("filter", "INPUT", "-j DIPLONAT", 1)?;
+
     Ok(())
 }
 
-pub fn open_ports(ipt: &iptables::IPTables, ports: Vec<Port>) -> Result<(), FirewallError> {
+pub fn open_ports(ipt: &iptables::IPTables, ports: messages::PublicExposedPorts) -> Result<(), FirewallError> {
+
+    for p in ports.tcp_ports {
+        ipt.append("filter", "DIPLONAT", &format!("-p tcp --dport {} -j ACCEPT", p))?;
+    }
 
-    for p in ports {
-        ipt.append("filter", "DIPLONAT", &format!("-p {} --dport {} -j ACCEPT", p.proto, p.number))?;
+    for p in ports.udp_ports {
+        ipt.append("filter", "DIPLONAT", &format!("-p udp --dport {} -j ACCEPT", p))?;
     }
 
     Ok(())
 }
 
-pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<HashSet<Port>, FirewallError> {
-    let mut opened_ports: HashSet<Port> = HashSet::new();
+pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExposedPorts, FirewallError> {
+    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();
@@ -50,13 +52,18 @@ pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<HashSet<Port>, Firew
                 let proto = String::from(raw_proto.as_str());
                 let number = String::from(raw_port.as_str()).parse::<u16>().unwrap();
 
-                opened_ports.insert( Port { proto, number } );
+                if proto == "tcp" {
+                    ports.tcp_ports.insert(number);
+                } else {
+                    ports.udp_ports.insert(number);
+                }
+
             },
             _ => {}
         }
     }
 
-    Ok(opened_ports)
+    Ok(ports)
 }
 
 pub fn cleanup(ipt: &iptables::IPTables) -> Result<(), FirewallError> {