add actor for firewall & massive refactor

1 files changed, 80 insertions, 0 deletions

diff --git a/src/fw_actor.rs b/src/fw_actor.rs
new file mode 100644
index 0000000..9bc6610
--- /dev/null
+++ b/src/fw_actor.rs
@@ -0,0 +1,80 @@
+use igd::aio::*;
+use igd::PortMappingProtocol;
+use std::net::SocketAddrV4;
+use log::*;
+use anyhow::{Result, Context};
+use tokio::{
+  select, 
+  sync::watch, 
+  time::{
+    self, 
+    Duration
+}};
+
+use iptables;
+use crate::messages;
+use crate::fw;
+use std::collections::HashSet;
+
+pub struct FirewallActor {
+  ipt: iptables::IPTables,
+  rx_ports: watch::Receiver<messages::PublicExposedPorts>,
+  last_ports: messages::PublicExposedPorts,
+  refresh: Duration
+}
+
+impl FirewallActor {
+  pub async fn new(_refresh: Duration, rxp: &watch::Receiver<messages::PublicExposedPorts>) -> Result<Self> {
+
+    
+    let ctx = Self { 
+      ipt: iptables::new(false).unwrap(),
+      rx_ports: rxp.clone(),
+      last_ports: messages::PublicExposedPorts::new(),
+      refresh: _refresh,
+    };
+
+    fw::setup(&ctx.ipt).expect("iptables setup failed");
+    
+    return Ok(ctx);
+  }
+
+  pub async fn listen(&mut self) -> Result<()> {
+    let mut interval = time::interval(self.refresh);
+    loop {
+      // 1. Wait for an event
+      let new_ports = select! {
+        Some(ports) = self.rx_ports.recv() => Some(ports),
+        _ = interval.tick() => None,
+        else => return Ok(()) // Sender dropped, terminate loop.
+      };
+
+      // 2. Update last ports if needed
+      if let Some(p) = new_ports { self.last_ports = p; }
+
+      // 3. Update firewall rules
+      match self.do_fw_update().await {
+        Ok(()) => debug!("Successfully updated firewall rules"),
+        Err(e) => error!("An error occured while updating firewall rules. {}", e),
+      }
+    }
+  }
+
+  pub async fn do_fw_update(&self) -> Result<()> {
+    
+    let curr_opened_ports = fw::get_opened_ports(&self.ipt).unwrap();
+
+    let diff_tcp = self.last_ports.tcp_ports.difference(&curr_opened_ports.tcp_ports).copied().collect::<HashSet<u16>>();
+    let diff_udp = self.last_ports.udp_ports.difference(&curr_opened_ports.udp_ports).copied().collect::<HashSet<u16>>();
+
+    let ports_to_open = messages::PublicExposedPorts {
+        tcp_ports: diff_tcp,
+        udp_ports: diff_udp
+    };
+
+    fw::open_ports(&self.ipt, ports_to_open).unwrap();
+
+    return Ok(());
+  }
+
+}