firewall/main.py

#! /usr/bin/python3

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

# Copyright © 2018-2019 Hugo Levy-Falk <hugo@klafyvel.me>

import os
import tempfile
import logging
from logging.handlers import RotatingFileHandler

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__))

LOG_LEVEL = logging.INFO

logger = logging.getLogger()
logger.setLevel(LOG_LEVEL)
formatter = logging.Formatter('%(levelname)s :: %(message)s')
file_handler = RotatingFileHandler('/var/log/firewall.log', 'a', 1000000, 1)
file_handler.setLevel(LOG_LEVEL)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
stream_handler = logging.StreamHandler()
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
@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("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__':
    cli()