Nicer cli.

This commit is contained in:
Hugo Levy-Falk 2019-05-06 23:01:59 +02:00 committed by root
parent 99bae1a02d
commit eac5407780
2 changed files with 89 additions and 10 deletions

94
main.py
View file

@ -1,5 +1,6 @@
#! /usr/bin/python3 #! /usr/bin/python3
import os import os
import tempfile
import logging import logging
from logging.handlers import RotatingFileHandler from logging.handlers import RotatingFileHandler
@ -7,6 +8,7 @@ import click
import nat as _nat import nat as _nat
import mac_ip as _mac_ip import mac_ip as _mac_ip
from firewall import CommandExec, ExecError
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) BASE_DIR = os.path.dirname(os.path.abspath(__file__))
@ -24,25 +26,99 @@ stream_handler.setFormatter(formatter)
stream_handler.setLevel(LOG_LEVEL) stream_handler.setLevel(LOG_LEVEL)
logger.addHandler(stream_handler) 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.group(invoke_without_command=True)
@click.pass_context @click.pass_context
def cli(ctx): @click.option('--keep-nat/--dont-keep-nat', default=False, help='Should I keep the current NAT table ?')
logger.info("Starting Re2o firewall manager.") @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: if ctx.invoked_subcommand is None:
logger.info("Loading firewall.") logger.info("Starting Re2o firewall manager.")
os.system('nft -I {install_dir} -f {firewall}'.format( _structure(keep_nat, keep_macip)
install_dir=BASE_DIR,
firewall=os.path.join(BASE_DIR, 'firewall.nft')
))
_mac_ip.update_macip()
_nat.main()
@cli.command() @cli.command()
def macip(): 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() _mac_ip.update_macip()
@cli.command() @cli.command()
def nat(): 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() _nat.main()
if __name__ == '__main__': if __name__ == '__main__':

5
nat.py
View file

@ -21,14 +21,17 @@ Creates the nat set.
import logging import logging
import time import time
import os
from configparser import ConfigParser from configparser import ConfigParser
import netaddr import netaddr
from firewall import NAT from firewall import NAT
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG = ConfigParser() CONFIG = ConfigParser()
CONFIG.read('/usr/local/firewall/config.ini') CONFIG.read(os.path.join(BASE_DIR, 'config.ini'))
def create_nat_adherent(): def create_nat_adherent():