From 578d5aa51a5f0cc254571a9ab75bd49be5388482 Mon Sep 17 00:00:00 2001
From: Hugo Levy-Falk <me@klafyvel.me>
Date: Mon, 6 May 2019 21:46:13 +0200
Subject: [PATCH] Nat ICMP

---
 firewall.py | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/firewall.py b/firewall.py
index 8313140..97f5c87 100755
--- a/firewall.py
+++ b/firewall.py
@@ -607,6 +607,12 @@ def get_ip_iterable_from_str(ip):
 
 class NAT:
 
+    PROTOCOLS = (
+            'tcp',
+            'udp',
+            'icmp'
+    )
+
     def __init__(self,
                  name,
                  range_in,
@@ -648,22 +654,16 @@ class NAT:
             grp: The name of the group
             ports: The port range (str)
         """
-        CommandExec.run([
-            *self.nft,
-            "add rule ip nat {name}_nat ip saddr @{name}_nat_port_{grp} ip protocol tcp snat ip saddr map @{name}_nat_address : {ports}".format(
-                name=self.name,
-                grp=grp,
-                ports=ports
-            )
-        ])
-        CommandExec.run([
-            *self.nft,
-            "add rule ip nat {name}_nat ip saddr @{name}_nat_port_{grp} ip protocol udp snat ip saddr map @{name}_nat_address : {ports}".format(
-                name=self.name,
-                grp=grp,
-                ports=ports
-            )
-        ])
+        for protocol in self.PROTOCOLS:
+            CommandExec.run([
+                *self.nft,
+                "add rule ip nat {name}_nat ip saddr @{name}_nat_port_{grp} ip protocol {protocol} snat ip saddr map @{name}_nat_address : {ports}".format(
+                    protocol=protocol,
+                    name=self.name,
+                    grp=grp,
+                    ports=ports
+                )
+            ])
 
     def manage(self):
         """Creates the port sets, ip map and rules