From eac54077801f08e7abea9e0bbfd749e653fe0f77 Mon Sep 17 00:00:00 2001
From: Hugo Levy-Falk <me@klafyvel.me>
Date: Mon, 6 May 2019 23:01:59 +0200
Subject: [PATCH] Nicer cli.

---
 main.py | 94 +++++++++++++++++++++++++++++++++++++++++++++++++++------
 nat.py  |  5 ++-
 2 files changed, 89 insertions(+), 10 deletions(-)

diff --git a/main.py b/main.py
index 987f114..fce53a9 100755
--- a/main.py
+++ b/main.py
@@ -1,5 +1,6 @@
 #! /usr/bin/python3
 import os
+import tempfile
 import logging
 from logging.handlers import RotatingFileHandler
 
@@ -7,6 +8,7 @@ import click
 
 import nat as _nat
 import mac_ip as _mac_ip
+from firewall import CommandExec, ExecError
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -24,25 +26,99 @@ stream_handler.setFormatter(formatter)
 stream_handler.setLevel(LOG_LEVEL)
 logger.addHandler(stream_handler)
 
+def _structure(keep_nat, keep_macip):
+    logger.info("Loading firewall.")
+    if keep_nat:
+        logging.info("Backing up the current NAT table.")
+        nat_file = tempfile.NamedTemporaryFile()
+        try:
+            code, nat, *_ = CommandExec.run_check_output([
+                'sudo',
+                '/usr/sbin/nft',
+                'list table nat'
+            ])
+        except ExecError as e:
+            logging.error(e)
+            return
+        else:
+            nat_file.write(nat.encode('utf-8'))
+    if keep_macip:
+        logging.info("Backing up the current macip set.")
+        macip_file = tempfile.NamedTemporaryFile()
+        try:
+            code, nat, *_ = CommandExec.run_check_output([
+                'sudo',
+                '/usr/sbin/nft',
+                'list set inet firewall ip_mac'
+            ])
+        except ExecError as e:
+            logging.error(e)
+            return
+        else:
+            macip_file.write(nat.encode('utf-8'))
+
+    CommandExec.run([
+        'nft', 
+        '-I',
+        BASE_DIR,
+        '-f',
+        os.path.join(BASE_DIR, 'firewall.nft')
+    ])
+    if keep_macip:
+        logging.info("Retreiving the current macip set.")
+        CommandExec.run([
+            'nft', 
+            '-I',
+            BASE_DIR,
+            '-f',
+            macip_file.name
+        ])
+        macip_file.close()
+    else:
+        _mac_ip.update_macip()
+    if keep_nat:
+        logging.info("Retreiving the current NAT table.")
+        CommandExec.run([
+            'nft', 
+            '-I',
+            BASE_DIR,
+            '-f',
+            nat_file.name
+        ])
+        nat_file.close()
+    else:
+        _nat.main()
+
+
 @click.group(invoke_without_command=True)
 @click.pass_context
-def cli(ctx):
-    logger.info("Starting Re2o firewall manager.")
+@click.option('--keep-nat/--dont-keep-nat', default=False, help='Should I keep the current NAT table ?')
+@click.option('--keep-macip/--dont-keep-macip', default=False, help='Should I keep the current macip set ?')
+def cli(ctx, keep_nat, keep_macip):
+    """Re2o firewall manager.
+
+    Used without command, the firewall manager will load the whole firewall (i.e. the struture, the macip set and the MAC table). By default it erases the current NAT table and macp set. You can choose to keep the current values for these with the flags.
+ 
+    """
     if ctx.invoked_subcommand is None:
-        logger.info("Loading firewall.")
-        os.system('nft -I {install_dir} -f {firewall}'.format(
-            install_dir=BASE_DIR,
-            firewall=os.path.join(BASE_DIR, 'firewall.nft')
-        ))
-        _mac_ip.update_macip()
-        _nat.main()
+        logger.info("Starting Re2o firewall manager.")
+        _structure(keep_nat, keep_macip)
 
 @cli.command()
 def macip():
+    """Load the macip set.
+
+    Load the macip set from Re2o. This mean you need to be able to contact the Re2o server :)
+    """
     _mac_ip.update_macip()
 
+    
 @cli.command()
 def nat():
+    """Load the NAT table.
+
+    Generate the NAT table from the config file. You typically need to run this command only at boot.
+    """
     _nat.main()
 
 if __name__ == '__main__':
diff --git a/nat.py b/nat.py
index 8604c1d..2af9011 100644
--- a/nat.py
+++ b/nat.py
@@ -21,14 +21,17 @@ Creates the nat set.
 
 import logging
 import time
+import os
 from configparser import ConfigParser
 
 import netaddr
 
 from firewall import NAT
 
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
 CONFIG = ConfigParser()
-CONFIG.read('/usr/local/firewall/config.ini')
+CONFIG.read(os.path.join(BASE_DIR, 'config.ini'))
 
 
 def create_nat_adherent():