#! /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 © 2019 Hugo Levy-Falk <me@klafyvel.me>

"""
Creates the nat set.
"""

import logging
from configparser import ConfigParser

import netaddr

from firewall import NetfilterSet

CONFIG = ConfigParser()
CONFIG.read('config.ini')


def get_ip_iterable_from_str(ip):
    try:
        ret = netaddr.IPGlob(ip)
    except netaddr.core.AddrFormatError:
        try:
            ret = netaddr.IPNetwork(ip)
        except netaddr.core.AddrFormatError:
            begin,end = ip.split('-')
            ret = netaddr.IPRange(begin,end)
    return ret


def create_nat(name, range_in, range_out, first_port, last_port):
    """Create two nftables tables for the nat:
        - <name>_address : which link a (or a range of) local address to a
        public address;
        - <name>_port : which links a local address to a
        range of ports.

        Args:
            name: name of the sets
            range_in: an IPRange with the private IP address
            range_out: an IPRange with the public IP address
            first_port: the first port used for the nat
            last_port: the last port used for the nat
        Returns:
            (<name>_address, <name>_port) which are NetfilterSet
    """
    assert 0 <= first_port < last_port < 65536, (name + ": Your first_port "
                                     "is lower than your last_port")
    nb_private_by_public = range_in.size // range_out.size
    nb_port_by_ip = (last_port - first_port + 1) // nb_private_by_public

    ports = {}
    ips = {}

    port = first_port
    for ip in range_in:
        ports[(str(netaddr.IPAddress(ip)),)] = ("%d-%d" % (port, min(port+nb_port_by_ip, 65535)),)
        port += nb_port_by_ip + 1
        if port >= last_port:
            port = first_port
    ip = range_in.first
    for ip_out in range_out:
        ips[('-'.join([
                str(netaddr.IPAddress(ip)),
                str(netaddr.IPAddress(ip+nb_private_by_public))
                ]),)] = (str(ip_out),)
        ip += nb_private_by_public + 1

    return (
        NetfilterSet(
            target_content=ips,
            type_=('IPv4',),
            name=name+'_nat_address',
            table_name='nat',
            flags=('interval',),
            type_from=('IPv4',)
        ),
        NetfilterSet(
            target_content=ports,
            type_=('port',),
            name=name+'_nat_port',
            table_name='nat',
            flags=('interval',),
            type_from=('IPv4',)
        ),
    )


def create_nat_adherent():
    range_in = get_ip_iterable_from_str(CONFIG['NAT']['range_in_adherent'])
    range_out = get_ip_iterable_from_str(CONFIG['NAT']['range_out_adherent'])
    first_port = int(CONFIG['NAT']['first_port_adherent'])
    last_port = int(CONFIG['NAT']['last_port_adherent'])
    return create_nat(
        'adherent',
        range_in,
        range_out,
        first_port,
        last_port
    )


def create_nat_federez():
    range_in = get_ip_iterable_from_str(CONFIG['NAT']['range_in_federez'])
    range_out = get_ip_iterable_from_str(CONFIG['NAT']['range_out_federez'])
    first_port = CONFIG['NAT']['first_port_federez']
    last_port = CONFIG['NAT']['last_port_federez']
    return create_nat(
        'federez',
        range_in,
        range_out,
        first_port,
        last_port
    )


def create_nat_aloes():
    range_in = get_ip_iterable_from_str(CONFIG['NAT']['range_in_aloes'])
    range_out = get_ip_iterable_from_str(CONFIG['NAT']['range_out_aloes'])
    first_port = CONFIG['NAT']['first_port_aloes']
    last_port = CONFIG['NAT']['last_port_aloes']
    return create_nat(
        'aloes',
        range_in,
        range_out,
        first_port,
        last_port
    )


def create_nat_admin():
    range_in = get_ip_iterable_from_str(CONFIG['NAT']['range_in_admin'])
    range_out = get_ip_iterable_from_str(CONFIG['NAT']['range_out_admin'])
    first_port = CONFIG['NAT']['first_port_admin']
    last_port = CONFIG['NAT']['last_port_admin']
    return create_nat(
        'admin',
        range_in,
        range_out,
        first_port,
        last_port
    )


def create_nat_prerezotage():
    range_in = get_ip_iterable_from_str(CONFIG['NAT']['range_in_prerezotage'])
    range_out = get_ip_iterable_from_str(CONFIG['NAT']['range_out_prerezotage'])
    first_port = CONFIG['NAT']['first_port_prerezotage']
    last_port = CONFIG['NAT']['last_port_prerezotage']
    return create_nat(
        'prerezotage',
        range_in,
        range_out,
        first_port,
        last_port
    )


def main():
    logging.info("Creating adherent nat...")
    address, port = create_nat_adherent()
    address.manage()
    port.manage()
    logging.info("Done.")
    #logging.info("Creating federez nat...")
    #address, port = create_nat_federez()
    #address.manage()
    #port.manage()
    #logging.info("Done.")
    #logging.info("Creating aloes nat...")
    #address, port = create_nat_aloes()
    #address.manage()
    #port.manage()
    #logging.info("Done.")
    #logging.info("Creating admin nat...")
    #address, port = create_nat_admin()
    #address.manage()
    #port.manage()
    #logging.info("Done.")
    #logging.info("Creating prerezotage nat...")
    #address, port = create_nat_prerezotage()
    #address.manage()
    #port.manage()
    #logging.info("Done.")


if __name__=='__main__':
    logging.info('Updating the NAT table.')
    main()