#!/usr/bin/env python
# -*- coding: utf-8 -*-
import utils
import base

from utils import pretty_print, OK, anim
from base import dev

class firewall(base.firewall_routeur):

    def __init__(self):
        super(self.__class__, self).__init__()

        self.reloadable.update({
            'log_all' : self.log_all,
            'admin_vlan' : self.admin_vlan,
            'clamp_mss' : self.clamp_mss,
            'ingress_filtering' : self.ingress_filtering,
            'ssh_on_https' : self.ssh_on_https,
            'connexion_secours' : self.connexion_secours,
            'connexion_appartement' : self.connexion_appartement,
            'blacklist_soft' : self.blacklist_soft,
            'blacklist_upload' : self.blacklist_upload,
            'reseaux_non_routable' : self.reseaux_non_routable,
            'filtrage_ports' : self.filtrage_ports,
            'limitation_debit' : self.limitation_debit,
            'limit_ssh_connexion' : self.limit_ssh_connexion,
            'tunnel_6in4' : self.tunnel_6in4,
        })

        self.use_ipset.extend([self.blacklist_soft, self.blacklist_upload, self.reseaux_non_routable])
        self.use_tc.extend([self.limitation_debit])

        self.ipset['reseaux_non_routable'] = {
            'deny' : base.Ipset("RESEAUX-NON-ROUTABLE-DENY","nethash"),
            'allow' :  base.Ipset("RESEAUX-NON-ROUTABLE-ALLOW","nethash"),
        }

        self.ipset['blacklist'].update({
            'soft' : base.Ipset("BLACKLIST-SOFT","ipmap","--from 138.231.136.0 --to 138.231.151.255"),
            'upload' : base.Ipset("BLACKLIST-UPLOAD","ipmap","--from 138.231.136.0 --to 138.231.151.255"),
        })

    def blacklist_maj(self, ips):
        self.blacklist_hard_maj(ips)
        self.blacklist_soft_maj(ips)
        self.blacklist_upload_maj(ips)

    def blacklists(self, table=None, fill_ipset=False, apply=False):
        self.blacklist_hard(table=table, fill_ipset=fill_ipset, apply=apply)
        self.blacklist_soft(table=table, fill_ipset=fill_ipset, apply=apply)
        self.blacklist_upload(table=table, fill_ipset=fill_ipset, apply=apply)

    def raw_table(self):
        """Génère les règles pour la table ``raw`` et remplis les chaines de la table"""
        table = 'raw'

        chain = 'PREROUTING'
        self.add(table, chain, '-d 225.0.0.50 -j DROP')
        return


    def mangle_table(self):
        table = 'mangle'
        super(self.__class__, self).mangle_table()

        chain = 'PREROUTING'
        self.add(table, chain, '-j %s' % self.log_all(table))
        self.add(table, chain, '-j %s' % self.connexion_secours(table))
        self.add(table, chain, '-p tcp -j CONNMARK --restore-mark')

        chain = 'POSTROUTING'
        self.add(table, chain, '-j %s' % self.clamp_mss(table))
        self.add(table,chain, '-j %s' % self.limitation_debit(table, run_tc=True))
        self.add(table, chain, '-j %s' % self.blacklist_upload(table, fill_ipset=True))
        return

    def filter_table(self):
        table = 'filter'
        super(self.__class__, self).filter_table()

        mac_ip_chain = self.test_mac_ip()
        blacklist_hard_chain = self.blacklist_hard()
        blacklist_soft_chain = self.blacklist_soft(table, fill_ipset=True)

        chain = 'INPUT'
        self.flush(table, chain)
        self.add(table, chain, '-i lo -j ACCEPT')
        self.add(table, chain, '-p icmp -j ACCEPT')
        self.add(table, chain, '-m state --state RELATED,ESTABLISHED -j ACCEPT')
        self.add(table, chain, '-j %s' % blacklist_soft_chain)
        for net in base.config.NETs['all'] + base.config.NETs['adm'] + base.config.NETs['personnel-ens']:
            self.add(table, chain, '-s %s -j %s' % (net, mac_ip_chain))
        self.add(table, chain, '-j %s' % blacklist_hard_chain)

        chain = 'FORWARD'
        self.flush(table, chain)
        self.add(table, chain, '-i lo -j ACCEPT')
        self.add(table, chain, '-j %s' % self.reseaux_non_routable(table, fill_ipset=True))
        self.add(table, chain, '-p icmp -j ACCEPT')
        self.add(table, chain, '-j %s' % self.admin_vlan(table))
        self.add(table, chain, '-j %s' % blacklist_soft_chain)
        self.add(table, chain, '-i %s -j %s' % (dev['out'], blacklist_hard_chain))
        self.add(table, chain, '-o %s -j %s' % (dev['out'], blacklist_hard_chain))
        self.add(table, chain, '-m state --state RELATED,ESTABLISHED -j ACCEPT')
        self.add(table, chain, '-j %s' % self.tunnel_6in4(table))
        for net in base.config.NETs['all'] + base.config.NETs['adm'] + base.config.NETs['personnel-ens']:
            self.add(table, chain, '-s %s -j %s' % (net, mac_ip_chain))
        self.add(table, chain, '-j %s' % self.connexion_secours(table))
        self.add(table, chain, '-j %s' % self.connexion_appartement(table))
        self.add(table, chain, '-j %s' % self.ingress_filtering(table))
        self.add(table, chain, '-j %s' % self.limit_ssh_connexion(table))
        self.add(table, chain, '-i %s -j %s' % (dev['out'], self.filtrage_ports(table)))
        self.add(table, chain, '-o %s -j %s' % (dev['out'], self.filtrage_ports(table)))
        return

    def nat_table(self):
        table = 'nat'
        super(self.__class__, self).nat_table()

        chain = 'PREROUTING'
        self.add(table, chain, '-j %s' % self.ssh_on_https(table))
        self.add(table, chain, '-j %s' % self.connexion_secours(table))
        self.add(table, chain, '-j %s' % self.blacklist_soft(table))

        chain = 'POSTROUTING'
        self.add(table, chain, '-j %s' % self.connexion_appartement(table))
        return

    def tunnel_6in4(self, table=None, apply=False):
        chain = 'TUNNEL_IPV6'

        tunnels_ipv6 = [ ('216.66.84.42', '138.231.136.12'), ('216.66.84.42','138.231.136.164') ]

        if table == 'filter':
            pretty_print(table, chain)
            for sideA, sideB in tunnels_ipv6:
                self.add(table, chain, '--proto 41 -s %s -d %s -j ACCEPT' % (sideA, sideB))
                self.add(table, chain, '--proto 41 -s %s -d %s -j ACCEPT' % (sideB, sideA))
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def limit_ssh_connexion(self, table=None, apply=False):
        chain = 'LIMIT-SSH-CONNEXION'

        if table == 'filter':
            pretty_print(table, chain)
            self.add(table, chain, '-i %s -p tcp --dport ssh -m state --state NEW -m recent --name SSH --set' % dev['out'])
            self.add(table, chain, '-i %s -p tcp --dport ssh -m state --state NEW -m recent --name SSH --update --seconds 30 --hitcount 10 --rttl -j DROP' % dev['out'])
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def test_mac_ip(self, table=None, fill_ipset=False, apply=False):
        chain = super(self.__class__, self).test_mac_ip()

        if table == 'filter':
            for key in ['out', 'tun-ovh' ]:
                self.add(table, chain, '-i %s -j RETURN' % dev[key])

        return super(self.__class__, self).test_mac_ip(table, fill_ipset, apply)


    def log_all(self, table=None, apply=False):
        chain = 'LOG_ALL'

        if table == 'mangle':
            pretty_print(table, chain)
            for device in dev.values():
                self.add(table, chain, '-i %s -m state --state NEW -j LOG --log-prefix "LOG_ALL "' % device)
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def admin_vlan(self, table=None, apply=False):
        chain = 'VLAN-ADM'

        if table == 'filter':
            pretty_print(table, chain)
            for net in base.config.NETs['adm']:
                self.add(table, chain, '-o %s -s %s -j ACCEPT' % (dev['tun-ovh'], net))
                self.add(table, chain, '-i %s -d %s -j ACCEPT' % (dev['tun-ovh'], net))
                self.add(table, chain, '-d %s -j REJECT' % net)
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def qos(self, table=None, apply=False):
        return

    def clamp_mss(self, table=None, apply=False):
        """Force la MSS (Max Segment Size) TCP à rentrer dans la MTU (Max Transfert Unit)"""
        chain = 'CLAMP-MSS'
        if table == 'mangle':
            pretty_print(table, chain)

            self.add(table, chain, '-p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu')
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def ingress_filtering(self, table=None, apply=False):
        """Pour ne pas router les paquêtes n'appartenant pas à notre plage ip voulant sortir de notre réseau
        et empêcher certain type de spoof (cf http://travaux.ovh.net/?do=details&id=5183)"""
        chain = 'INGRESS_FILTERING'
        if table == 'filter':
            pretty_print(table, chain)

            for net in base.config.NETs['all']:
                self.add(table, chain, '-o %s -s %s -j RETURN' % (dev['out'], net))
            self.add(table, chain, '-o %s -j LOG --log-prefix BAD_ROUTE' % dev['out'])
            self.add(table, chain, '-o %s -j DROP' % dev['out'])
            for net_d in base.config.NETs['all']:
                for net_s in base.config.NETs['all']:
                    self.add(table, chain,'-i %s ! -s %s -d %s -j RETURN' % (dev['out'], net_s, net_d))
            self.add(table, chain,'-i %s -j LOG --log-prefix BAD_SRC' % dev['out'])
            self.add(table, chain,'-i %s -j DROP' % dev['out'])
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def ssh_on_https(self, table=None, apply=False):
        """Pour faire fonctionner ssh2.crans.org"""
        chain = 'SSH2'

        if table == 'nat':
            pretty_print(table, chain)
            self.add(table, chain, '-p tcp -d 138.231.136.2 --dport 22 -j DNAT --to-destination 138.231.136.1:22')    # redirection du ssh vers zamok
            self.add(table, chain, '-p tcp -d 138.231.136.2 --dport 443 -j DNAT --to-destination 138.231.136.1:22')    # redirection du ssh vers zamok (pour passer dans un proxy, avec corkscrew)
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def connexion_secours(self, table=None, apply=False):
        """Redirige les paquets vers un proxy lorsqu'on est en connexion de secours"""
        chain = 'CONNEXION-SECOURS'

        if table == 'mangle':
            pretty_print(table, chain)
            self.add(table, chain, '-p tcp -s 138.231.136.0/16 ! -d 138.231.136.0/16 --destination-port 80 -m condition --condition secours -j MARK --set-mark %s' % (base.config.firewall.mark['secours']))
            self.add(table, chain, '-m mark --mark %s -j ACCEPT' % base.config.firewall.mark['secours'])
            print OK

        if table == 'nat':
            pretty_print(table, chain)
            self.add(table, chain, '-p tcp -m mark --mark %s -j DNAT --to-destination 10.231.136.4:3129' % base.config.firewall.mark['secours'] )
            print OK

        if table == 'filter':
            pretty_print(table, chain)
            self.add(table, chain, '-p tcp -s 138.231.136.0/16 ! -d 138.231.136.0/16 --destination-port 443 -m condition --condition secours -j REJECT')
            self.add(table, chain, '-m mark --mark %s -j ACCEPT' % base.config.firewall.mark['secours'])
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def connexion_appartement(self, table=None, apply=False):
        """PNAT les appartements derrière appartement.crans.org"""
        chain = 'CONNEXION-APPARTEMENT'

        if table == 'nat':
            pretty_print(table, chain)
            for dev_key in ['out', 'fil', 'wifi']:
                for net in base.config.NETs['personnel-ens']:
                    self.add(table, chain, '-o %s -s %s -j SNAT --to 138.231.136.44' % (dev[dev_key], net))
            print OK

        if table == 'filter':
            pretty_print(table, chain)
            for net in base.config.NETs['personnel-ens']:
                self.add(table, chain, '-s %s -j ACCEPT' % net)
                self.add(table, chain, '-d %s -j ACCEPT' % net)
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def blacklist_soft_maj(self, ip_list):
        self.blacklist_soft(fill_ipset=True)
#        for ip in ip_list:
#            machine = self.conn.search(u"ipHostNumber=%s" % ip)
#            # Est-ce qu'il y a des blacklists soft parmis les blacklists de la machine
#            if machine and set([bl['type'] for bl in machine[0].blacklist_actif() ]).intersection(base.config.blacklist_sanctions_soft):
#                try: self.ipset['blacklist']['soft'].add(ip)
#                except IpsetError: pass
#            else:
#                try: self.ipset['blacklist']['soft'].delete(ip)
#                except IpsetError: pass

    def blacklist_soft(self, table=None, fill_ipset=False, apply=False):
        """Redirige les gens blacklisté vers le portail captif"""
        chain = 'BLACKLIST_SOFT'

        if fill_ipset:
            # On récupère la liste de toutes les ips blacklistés soft
            bl_soft_ips = self.blacklisted_ips(base.config.blacklist_sanctions_soft, base.config.NETs['all'])
            anim('\tRestoration de l\'ipset %s' % self.ipset['blacklist']['soft'])
            self.ipset['blacklist']['soft'].restore(bl_soft_ips)
            print OK

        if table == 'filter':
            pretty_print(table, chain)
            self.add(table, chain, '-p tcp --dport 80 -m set --match-set %s src -j ACCEPT' % self.ipset['blacklist']['soft'] )
            self.add(table, chain, '-p tcp --sport 80 -m set --match-set %s dst -j ACCEPT' % self.ipset['blacklist']['soft'] )
            self.add(table, chain, '-p tcp -d 10.231.136.4 --dport 3128 -m set --match-set %s src -j ACCEPT' % self.ipset['blacklist']['soft'] )
            self.add(table, chain, '-p tcp -s 10.231.136.4 --sport 3128 -m set --match-set %s dst -j ACCEPT' % self.ipset['blacklist']['soft'] )
            print OK

        if table == 'nat':
            pretty_print(table, chain)
            for net in base.config.NETs['all']:
                self.add(table, chain, '-d %s -j RETURN' % net)
            self.add(table, chain, '-p tcp --dport 80 -m set --match-set %s src -j DNAT --to-destination 10.231.136.4:3128' % self.ipset['blacklist']['soft'] )
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def blacklist_upload_maj(self, ip_list):
        self.blacklist_upload(fill_ipset=True)
#        for ip in ip_list:
#            machine = self.conn.search(u"ipHostNumber=%s" % ip)
#            # Est-ce qu'il y a des blacklists pour upload parmis les blacklists de la machine
#            if machine and set([bl['type'] for bl in machine[0].blacklist_actif() ]).intersection(blacklist_bridage_upload):
#                try: self.ipset['blacklist']['upload'].add(ip)
#                except IpsetError: pass
#            else:
#                try: self.ipset['blacklist']['upload'].delete(ip)
#                except IpsetError: pass

    def blacklist_upload(self, table=None, fill_ipset=False, apply=False):
        """Les gens blacklistés ne sont plus prioritaires (classe de qos commune)"""
        chain = 'BLACKLIST_UPLOAD'

        if fill_ipset:
            # On récupère la liste de toutes les ips blacklistés pour upload
            bl_upload_ips = self.blacklisted_ips(base.config.blacklist_bridage_upload, base.config.NETs['all'])
            anim('\tRestoration de l\'ipset %s' % self.ipset['blacklist']['upload'])
            self.ipset['blacklist']['upload'].restore(bl_upload_ips)
            print OK

        if table == 'mangle':
            pretty_print(table, chain)
            # Classification pour les blacklists upload
            self.add(table, chain, '-o %s -m set --match-set %s src -j CLASSIFY --set-class 1:11' % (dev['out'], self.ipset['blacklist']['upload']))

            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def reseaux_non_routable(self, table=None, fill_ipset=False, apply=False):
        """Bloque les réseaux non routables autres que ceux utilisés par le crans"""
        chain = 'RESEAUX_NON_ROUTABLES'

        if fill_ipset:
            anim('\tRestoration de l\'ipset reseaux_non_routable')
            allowed = [ net for nets in base.config.NETs.values() for net in nets if utils.NetInNets(net, base.config.firewall.reseaux_non_routables) ]
            self.ipset['reseaux_non_routable']['allow'].restore(allowed)
            self.ipset['reseaux_non_routable']['deny'].restore(base.config.firewall.reseaux_non_routables)
            print OK

        if table == 'filter':
            pretty_print(table, chain)
            self.add(table, chain, '-m set --match-set %s src -j RETURN' %  self.ipset['reseaux_non_routable']['allow'])
            self.add(table, chain, '-m set --match-set %s dst -j RETURN' %  self.ipset['reseaux_non_routable']['allow'])
            self.add(table, chain, '-m set --match-set %s src -j DROP' %  self.ipset['reseaux_non_routable']['deny'])
            self.add(table, chain, '-m set --match-set %s dst -j DROP' %  self.ipset['reseaux_non_routable']['deny'])
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def filtrage_ports_maj(self, ip_lists):
        self.filtrage_ports('filter', apply=True)

    def filtrage_ports(self, table=None, apply=False):
        """Ouvre les ports vers et depuis les machines du réseau crans"""
        chain = 'FILTRAGE-PORTS'

        def format_port(port):
            port = str(port)
            if port.endswith(':'):
                port = '%s65535' % port
            if port.startswith(':'):
                port = '0%s' % port
            return port

        def add_ports(ip, machine, proto, sens):
            self.add(
                table,
                chain,
                '-p %s -%s %s -m multiport --dports %s -j RETURN' % (
                        proto,
                        (sens=='out' and 's') or (sens == 'in' and 'd'),
                        ip,
                        ','.join( format_port(port) for port in machine['port%s%s' % (proto.upper(), sens)])
                )
            )

        if table == 'filter':
            pretty_print(table, chain)
            for net in base.config.NETs['adherents'] + base.config.NETs['wifi-adh'] + base.config.NETs['personnel-ens']:
                for proto in base.config.firewall.ports_default.keys():
                    if base.config.firewall.ports_default[proto]['output']:
                        self.add(table, chain, '-p %s -s %s -m multiport --dports %s -j RETURN' % (proto, net, ','.join( format_port(port) for port in base.config.firewall.ports_default[proto]['output'])))
                    if base.config.firewall.ports_default[proto]['input']:
                        self.add(table, chain, '-p %s -d %s -m multiport --dports %s -j RETURN' % (proto, net, ','.join( format_port(port) for port in base.config.firewall.ports_default[proto]['input'])))

            for machine in self.machines():
                for ip in machine['ipHostNumber']:
                    if 'portTCPout' in machine:
                        add_ports(ip, machine, 'tcp', 'out')
                    if 'portUDPout' in machine:
                        add_ports(ip, machine, 'udp', 'out')
                    if 'portTCPin' in machine:
                        add_ports(ip, machine, 'tcp', 'in')
                    if 'portUDPin' in machine:
                        add_ports(ip, machine, 'udp', 'in')

            self.add(table, chain, '-j REJECT')
            print OK

        if apply:
            self.apply(table, chain)
        return chain

    def limitation_debit(self, table=None, run_tc=False, apply=False):
        """Limite le débit de la connexion selon l'agréement avec l'ENS"""
        chain = 'LIMITATION-DEBIT'

        debit_max = base.config.firewall.debit_max
        uplink_speed = '1024mbit'

        if table == 'mangle':
            pretty_print(table, chain)
            # Pas de QoS vers/depuis la zone ENS
            self.add(table, chain, '-d 138.231.0.0/16 -s 138.231.0.0/16 -j RETURN')

            # Idem pour le ftp
            self.add(table, chain, '-d 138.231.136.98 -j RETURN')
            self.add(table, chain, '-s 138.231.136.98 -j RETURN')

            # Idem vers OVH pour le test de la connection de secours
            self.add(table, chain, '-d 91.121.84.138 -j RETURN')
            self.add(table, chain, '-s 91.121.84.138 -j RETURN')

            # Classification par defaut pour tous les paquets
            for net in base.config.NETs['all']:
                self.add(table, chain, '-o %s -s %s -j CLASSIFY --set-class 1:10' % (dev['out'], net))
                self.add(table, chain, '-o %s -d %s -j CLASSIFY --set-class 1:10' % (dev['fil'], net))
                self.add(table, chain, '-o %s -d %s -j CLASSIFY --set-class 1:10' % (dev['wifi'], net))

            # Classification pour les appartements
            for net in base.config.NETs['personnel-ens']:
                self.add(table, chain, '-o %s -d %s -j CLASSIFY --set-class 1:3' % (dev['app'], net))
                self.add(table, chain, '-o %s -s %s -j CLASSIFY --set-class 1:2' % (dev['out'], net))

            # Classification pour la voip
            self.add(table, chain, '-d sip.crans.org -j CLASSIFY --set-class 1:12')
            self.add(table, chain, '-s sip.crans.org -j CLASSIFY --set-class 1:12')
            print OK

        if run_tc:
            anim('\tApplication des commandes tc')
            for int_key in ['out', 'fil', 'wifi']:
                try:
                    utils.tc('qdisc del dev %s root' % dev[int_key])
                except utils.TcError:
                    pass
                utils.tc('qdisc add dev %s root handle 1: htb r2q 1' % dev[int_key])
                utils.tc("class add dev %s parent 1: classid 1:1 "
                   "htb rate %s ceil %s" % (dev[int_key], uplink_speed, uplink_speed))
                utils.tc("class add dev %s parent 1:1 classid 1:2 "
                   "htb rate %skbps ceil %skbps" % (dev[int_key], debit_max, debit_max))

                # Classe par defaut
                utils.tc('class add dev %s parent 1:2 classid 1:10 '
                       'htb rate %skbps ceil %skbps prio 1' % (dev[int_key], debit_max, debit_max))
                utils.tc('qdisc add dev %s parent 1:10 '
                       'handle 10: sfq perturb 10' % dev[int_key])

                # Classe par pour la voip
                utils.tc('class add dev %s parent 1:2 classid 1:12 '
                       'htb rate %skbps ceil %skbps prio 0' % (dev[int_key], debit_max, debit_max))
                utils.tc('qdisc add dev %s parent 1:12 '
                       'handle 12: sfq perturb 10' % dev[int_key])

            #Classe des decos upload
            utils.tc('class add dev %s parent 1:2 classid 1:11 '
                    'htb rate 60kbps ceil 60kbps prio 1' % dev['out'])
            utils.tc('qdisc add dev %s parent 1:11 '
                    'handle 11: sfq perturb 10' % dev['out'])

            for int_key in ['app']:
                try:
                    utils.tc('qdisc del dev %s root' % dev[int_key])
                except TcError:
                    pass
                utils.tc('qdisc add dev %s root handle 1: htb r2q 1' % dev[int_key])

                utils.tc("class add dev %s parent 1: classid 1:1 "
                   "htb rate 128kbps ceil 128kbps" % dev[int_key])

                # Classe pour l'upload des appartements
                utils.tc("class add dev %s parent 1:1 classid 1:2 "
                   "htb rate 128kbps ceil 128kbps" % dev[int_key])
                utils.tc('qdisc add dev %s parent 1:2 '
                   'handle 2: sfq perturb 10' % dev[int_key])

                # Classe pour le download des apparetments
                utils.tc("class add dev %s parent 1: classid 1:3 "
                   "htb rate %skbps ceil %skbps" % (dev[int_key], debit_max/10, debit_max/2))
                utils.tc('qdisc add dev %s parent 1:3 '
                   'handle 3: sfq perturb 10' % dev[int_key])

            print OK

        if apply:
            self.apply(table, chain)
        return chain