#! /usr/bin/env python
# -*- coding: iso-8859-15 -*-

"""
Script de déconnection automatique des machines du crans pour les raisons :
    - upload
    - p2p
    - flood
    - virus

Copyright (C) Xavier Pessoles, Étienne Chové, Vincent Bernat, Nicolas Salles
Licence : GPL v2
"""

################################################################################
#                          Import des commandes                                #
################################################################################

import commands
import sys
import psycopg2
sys.path.append('/usr/scripts/gestion')
from config import upload, virus, p2p, NETs
import ipt
import smtplib
from ldap_crans import crans_ldap
from ldap_crans import MachineWifi
from ldap_crans import MachineFixe
from time import *
#import locale
#locale.setlocale(locale.LC_TIME, 'fr_FR')
import lock
sys.path.append('/usr/scripts/surveillance/fiche_deconnexion')
from generate import generate_ps
from affich_tools import tableau

# ldap
ldap = crans_ldap()

# Quelques fonctions
####################

def machine_online(machine):
    """
    Retourne True si la machine est connectée au réseau et False si elle ne l'est pas
    """
    # Les machines wifi sont toujours online
    if isinstance(machine, MachineWifi):
        return True
    # Arping pour les fixes
    return not commands.getstatusoutput('/usr/sbin/arping -c 3 %s' % machine.mac())[0]

# Variables utiles
##################

# Création d'une chaine qui ressemble à :
# (ip_src<<=inet('138.231.136.0/21') or ip_src<<=inet('138.231.148.0/22'))
ip_src_in_crans = '(%s)' % ' or '.join([ "ip_src<<=inet('%s')" % net for net in NETs['all'] ])

# Connections :
###############

# Connection à la base sql via pgsql
pgsql = psycopg2.connect(database='filtrage', user='crans')
# Il faudra remplacer la ligne ci-dessous par pgsql.set_session(autocommit = True) sous wheezy
pgsql.set_isolation_level(0)
curseur = pgsql.cursor()

# Le smtp est assez capricieux
def connectsmtp():
    for i in range(5):
        try:
            mail = smtplib.SMTP('localhost')
        except:
            sleep(5)
            if i == 4:
                print "Impossible de se connecter au SMTP"
                sys.exit(1)
    return mail

# Pour trouver la chambre où était la machine que l'on déconnecte.
def reperage_chambre(mac):
    pgsql = psycopg2.connect(database='mac_prises', user='crans')
    # Il faudra remplacer la ligne ci-dessous par pgsql.set_session(autocommit = True) sous wheezy
    curseur = pgsql.cursor()
    requete = "SELECT date, chambre FROM correspondance WHERE mac=%s ORDER BY date DESC LIMIT 1;"
    curseur.execute(requete, (mac,))
    result = curseur.fetchall()
    if result:
        return result[0][0], result[0][1]
    else:
        return "Inconnue", "Inconnue"

################################################################################
#                         Vérification de l'upload                             #
################################################################################

# upload par entité (adhérent/club/machine crans)
upload4="""SELECT
    'upload', sum(upload)/1024/1024 AS total, ip_crans
FROM
    upload
WHERE
        upload > download
    AND date > timestamp 'now' - interval '1 day'
    AND date < 'now'
    AND NOT EXISTS
    (
        SELECT 1
        FROM exemptes
        WHERE upload.ip_crans <<= exemptes.ip_crans
        AND upload.ip_ext <<= exemptes.ip_dest
    )
GROUP BY
    ip_crans
"""
upload6 = """SELECT 
    'upload', sum(upload)/1024/1024 AS total, ip_crans 
FROM
    (
        SELECT DISTINCT * FROM 
        (
            SELECT 
                upload6.date, mac_ip.ip AS ip_crans, upload6.ip_ext, upload6.id, upload6.port_crans, upload6.port_ext, upload6.download, upload6.upload
            FROM mac_ip,upload6 
            WHERE 
                upload6.ip_crans = mac_ip.ip 
                AND upload6.date > mac_ip.date 
                AND upload6.date - interval '1 day' < mac_ip.date 
                AND upload6.date > timestamp 'now' - interval '1 day'
                AND upload6.date < 'now'
                AND upload6.upload > upload6.download
                AND NOT EXISTS
                (
                    SELECT 1
                    FROM exemptes
                    WHERE upload6.ip_crans <<= exemptes.ip_crans
                    AND upload6.ip_ext <<= exemptes.ip_dest
                )
        ) AS upload
    ) AS upload
WHERE
    upload > download
GROUP BY
    ip_crans
"""
requete = """SELECT
  round(total) AS total, machines.type AS type, machines.id AS id
FROM
  (
    (%s) UNION (%s)
  )
AS
  upload
INNER JOIN
  machines
ON
  machines.ip = upload.ip_crans
WHERE
  total >= 250
GROUP BY
  total, type, id
;""" % (upload4, upload6)
curseur.execute(requete)
uploadeurs = curseur.fetchall()


# On regarde s'il y a deux ipv6 identiques avec des mac non identiques
collision_mac_ip_request = "SELECT DISTINCT (a.*) FROM mac_ip as a, mac_ip as b where a.ip=b.ip AND a.mac != b.mac AND a.date >= b.date AND a.date - b.date < interval '3 day' ORDER BY a.date;"
curseur.execute(collision_mac_ip_request)
collision_mac_ip = curseur.fetchall()

if collision_mac_ip != []:
    print "Collision d'addresses ipv6 : "
for (date, mac, ip) in collision_mac_ip:
    print "%s %s %s" % (date, ipt.mac_addr(mac), ip)

# Table des avertis
###################

# Avertis upload hard
requete = "SELECT type,id FROM avertis_upload_hard where date>timestamp 'now' - interval '1 day'"
curseur.execute(requete)
avertis_upload_hard = curseur.fetchall()

# Avertis upload soft
requete = "SELECT type,id FROM avertis_upload_soft where date>timestamp 'now' - interval '1 day'"
curseur.execute(requete)
avertis_upload_soft = curseur.fetchall()

# Vérification :
################

for elupload, eltype, elid in uploadeurs:
    if elupload >= upload.hard:
        # L'adhérent a t il été blacklisté ?
        ####################################
        if (eltype, elid) in avertis_upload_hard:
            continue

        # Propriétaire issu de LDAP
        ###########################
        if eltype == 'club':
            proprio = ldap.search('cid=%d'%elid, 'w')['club']
        elif eltype == 'adherent':
            proprio = ldap.search('aid=%d'%elid, 'w')['adherent']
        else:
            continue
        

        if len(proprio) != 1:
            print 'Erreur : Proprio non trouvé (%s) %d'%(eltype, elid)
            continue
        proprio = proprio[0]

        # On cherche à savoir où et quand on
        # a vu les machines du proprio pour la dernière fois
        ####################################################
        machines = proprio.machines()
        
        macs_dates_chambres = []
        for machine in machines:
            if isinstance(machine, MachineFixe):
               mac = machine.mac() 
               date, chambre = reperage_chambre(mac) 
               macs_dates_chambres.append([mac, date, chambre])

        mdcf = tableau(macs_dates_chambres, ('mac', 'date', 'chambre'), (20, 21, 7), ('c', 'c', 'c'))

        # On inscrit l'instance dans la table des avertis_hard
        ######################################################
        curseur.execute("INSERT INTO avertis_upload_hard (type,id,date) VALUES ('%s','%d','now')"%(eltype,elid))

        # On sanctionne
        ###############
        debut = int(time())
        fin = debut + 24*3600
        proprio.blacklist([debut, fin, 'autodisc_upload', "Déconn auto. %s Mo" % elupload])
        proprio.save()

        # On envoie un mail à l'adhérent
        ################################
        mail = connectsmtp()

        corps = upload.message_hard % {'from': upload.expediteur, 'to': proprio.email(), 'upload': elupload, 'proprio': proprio.Nom()}
        corps = corps.encode('utf-8')
        mail.sendmail(upload.expediteur, proprio.email(), corps)

        # On envoie un mail à disconnect
        ################################
        if upload.disconnect_mail_hard:
            corps = upload.message_disconnect_hard % {'from': upload.expediteur, 'to': upload.expediteur, 'upload': elupload, 'proprio': proprio.Nom(), 'mdc': mdcf, 'chambre': proprio.chbre()}
            corps = corps.encode('utf-8')
            mail.sendmail(upload.expediteur, upload.expediteur, corps)

        # Vérification du nombre de déconnexions
        #########################################
        nb_decos = len([ x for x in proprio.blacklist() if int(x.split('$')[0]) > time()-30*24*3600 and x.split('$')[2] == 'autodisc_upload' ])
        if nb_decos >= 3:

            # Génération du fichier postscript
            try:
                fichier_ps = generate_ps('upload', proprio, ldap)
            except:
                fichier_ps = ("ERREUR lors de la génération. Merci de regénérer manuellement la fiche avec la commande :\n"
                    + "/usr/scripts/surveillance/fiche_deconnexion/generate.py --upload aid=%d" % int(proprio.id()))

            # Envoi du mail à disconnect
            corps = upload.message_disconnect_multi % {'from': upload.expediteur, 'to': upload.expediteur, 'nbdeco': nb_decos, 'proprio': proprio.Nom(), 'ps': fichier_ps}
            corps = corps.encode('utf-8')
            mail.sendmail(upload.expediteur, upload.expediteur, corps)

    elif elupload >= upload.soft:
        # L'adhérent a t il été averti ou est déjà déco ?
        #################################################
        if (eltype, elid) in avertis_upload_soft or (eltype, elid) in avertis_upload_hard:
            continue

        # Objets LDAP
        #############
        if eltype == 'club':
            proprio = ldap.search('cid=%d'%elid)['club']
        elif eltype == 'adherent':
            proprio = ldap.search('aid=%d'%elid)['adherent']
        else:
            continue

        if len(proprio) != 1:
            print 'Proprio non trouvé (%s) %d'%(eltype, elid)
            continue
        proprio = proprio[0]

        # On inscrit l'ip dans la table des avertis soft
        ################################################
        curseur.execute("INSERT INTO avertis_upload_soft (type,id,date) VALUES ('%s','%d','now')"%(eltype, elid))

        # On envoie un mail à l'adhérent
        ################################
        mail = connectsmtp()

        corps = upload.message_soft % {'from': upload.expediteur, 'to': proprio.email(), 'upload': elupload, 'proprio': proprio.Nom()}
        corps = corps.encode('utf-8')
        mail.sendmail(upload.expediteur, proprio.email(), corps)

        # On envoie un mail à disconnect
        ################################
        if upload.disconnect_mail_soft:
            corps = upload.message_disconnect_soft % {'from': upload.expediteur, 'to': upload.expediteur, 'upload': elupload, 'proprio': proprio.Nom()}
            corps = corps.encode('utf-8')
            mail.sendmail(upload.expediteur, upload.expediteur, corps)

# On supprime les vieux avertisements
curseur.execute("DELETE FROM avertis_upload_hard WHERE date < timestamp 'now' - interval '1 day'")
curseur.execute("DELETE FROM avertis_upload_soft WHERE date < timestamp 'now' - interval '1 day'")

################################################################################
#                   Détection de l'existence de virus                          #
################################################################################

# Dans la table virus on sélectionne les ip_src qui appartiennent au réseau
requete = "SELECT ip_src,count(ip_src) FROM virus WHERE %s and date > timestamp 'now' - interval '1 hour' group by ip_src" % ip_src_in_crans
curseur.execute(requete)
infectes = curseur.fetchall()

# Récupération des infectés pour ne pas les reblacklister
requete = "SELECT ip_crans FROM avertis_virus"
curseur.execute(requete)
infectes_old = curseur.fetchall()

for ip, nombre in infectes:

    # Si on est en dessous du seuil, on laisse passer
    if nombre < virus.virus:
        continue

    # Si on est déja avertis, on laisse passer
    if (ip) in infectes_old:
        continue

    # Lecture des infos de ldap
    machine = ldap.search('ipHostNumber=%s' % ip, 'w' )['machine'][0]
    hostname = machine.nom()
    proprio = machine.proprietaire()
    blacklist = proprio.blacklist()

    # Inscription dans la table des infectés
    requete = "INSERT INTO avertis_virus (ip_crans,date) VALUES ('%s','now')" % ip
    curseur.execute(requete)

    # On récupère les index des lignes de bl où il y a marqué virus
    index = [blacklist.index(x) for x in blacklist if 'autodisc_virus' in x ]
    if index:
        # L'adhérent est déjà blacklisté
        proprio.blacklist((index[0], ['now', '-', 'autodisc_virus', hostname]))
        proprio.save()
    else:
        # L'adhérent n'est pas encore blacklisté
        proprio.blacklist(['now', '-', 'autodisc_virus', hostname])
        proprio.save()

################################################################################
#                   Détection des virus qui floodent                           #
################################################################################

# Dans la table virus on sélectionne les ip_src qui appartiennent au réseau
requete = "SELECT ip_src,count(ip_src) FROM flood WHERE %s and date > timestamp 'now' - interval '1 hour' GROUP BY ip_src" % ip_src_in_crans
curseur.execute(requete)
infectes = curseur.fetchall()

# Récupération des infectés pour ne pas les reblacklister
requete = "SELECT ip_crans FROM avertis_virus"
curseur.execute(requete)
infectes_old = curseur.fetchall()

for ip, nombre in infectes:

    # Si on est en dessous du seuil, ou qu'on est déjà averti, on laisse passer
    if nombre < virus.flood or (ip) in infectes_old:
        continue

    # Lecture des infos de ldap
    try :
        machine = ldap.search('ipHostNumber=%s' % ip, 'w' )['machine'][0]
    except IndexError :
        # Dans le cas où l'ip détectée n'est pas enregistrée
        print "La machine avec l'ip %s n'est pas declaree !" % ip
        continue
    hostname = machine.nom()
    proprio = machine.proprietaire()
    blacklist = proprio.blacklist()

    # Inscription dans la table des infectés
    requete = "INSERT INTO avertis_virus (ip_crans,date) VALUES ('%s','now')" % ip
    curseur.execute(requete)

    # On récupère les index des lignes de bl où il y a marqué virus
    index = [ blacklist.index(x) for x in blacklist if 'autodisc_virus' in x ]
    try:
        if index:
            # L'adhérent est déjà blacklisté
            proprio.blacklist((index[0], ['now', '-', 'autodisc_virus', hostname]))
            proprio.save()
        else:
            # L'adhérent n'est pas encore blacklisté
            proprio.blacklist(['now', '-', 'autodisc_virus', hostname])
            proprio.save()
    except ValueError: # On a essayé de blacklister un proporiétaire virtuel
        pass # Le message d'erreur a déjà été affiché (changer ça ?)


# Reconnexion si le virus/flood a disparu
#########################################

# Dans la table avertis_virus on récupère la liste des infectés
requete = "SELECT ip_crans FROM avertis_virus where date < timestamp 'now' - interval '1 hour'"
curseur.execute(requete)
infectes = [ x[0] for x in curseur.fetchall() ]

for IP in infectes:

    # Nombre de requêtes de virus
    requete1 = "SELECT COUNT(ip_src) FROM virus where ip_src='%s' and date > timestamp 'now' - interval '1 hour'" % IP
    curseur.execute(requete1)
    nb_virus = curseur.fetchall()

    # Nombre de requêtes de flood
    requete2 = "SELECT COUNT(ip_src) FROM flood where ip_src='%s' and date > timestamp 'now' - interval '1 hour'" % IP
    curseur.execute(requete2)
    nb_flood = curseur.fetchall()

    # On ne traite que les IP qui sont descendues en dessous des seuils
    if nb_virus[0][0] < virus.virus and nb_flood[0][0] < virus.flood:
        try:
            machine = ldap.search('ipHostNumber=%s' % IP, 'w' )['machine'][0]
        except IndexError:
            print "Suppression de %s des machines infectees (la machine n'existe plus)"%IP
            requete = "DELETE FROM avertis_virus where ip_crans='%s'"%IP
            curseur.execute(requete)
            continue # la machine n'existe plus, on passe à l'infecté suivant
        # Si la machine n'est pas online, on reconnecte
        #if machine_online(machine):
        proprio = machine.proprietaire()
        bl = proprio.blacklist()
        hostname = machine.nom()
        # On stoppe la sanction pour une ligne existante de la blackliste
        # En prenant en compte le fait que d'autres lignes de blackliste
        # ont pu s'ajouter.
        lignes_enlevees = 0
        for ligne in bl:
            if '$-$autodisc_virus$%s' % hostname in ligne:
                liste = ligne.split('$')
                argument = [liste[0], 'now', liste[2], liste[3]]
                index = bl.index(ligne)
                proprio.blacklist((index, argument))
                proprio.save()
                lignes_enlevees += 1
        if lignes_enlevees == 0:
            print "Suppression de %s des machines infectees, mais aucune blackliste"%hostname
        requete = "DELETE FROM avertis_virus where ip_crans='%s'"%IP
        curseur.execute(requete)

################################################################################
#                   Gestion du peer to peer                                    #
################################################################################

# Dans la table p2p on sélectionne les ip_src qui appartiennent au réseau
requete = "SELECT ip_src,id_p2p,count(ip_src) FROM p2p WHERE %s AND date > timestamp 'now' - interval '2 hours' AND id_p2p != 1 GROUP BY ip_src,id_p2p ORDER BY ip_src" % ip_src_in_crans
curseur.execute(requete)
fraudeurs = curseur.fetchall()

# Récupération des fraudeurs pour ne pas les resanctionner
requete = "SELECT ip_crans,protocole FROM avertis_p2p WHERE  date > timestamp 'now' - interval '3 hours'"
curseur.execute(requete)
avertisp2p = curseur.fetchall()

for ip, id_p2p, nombre in fraudeurs:

    # On récupére le protocole de p2p :
    requete = "SELECT nom FROM protocole_p2p WHERE id_p2p=%d" % id_p2p
    curseur.execute(requete)
    protocole = curseur.fetchall()[0][0]

    # On ne prend pas en compte s'il est sous le seuil admis, ou
    #s'il est averti
    if nombre <= p2p.limite[protocole] or (ip, protocole) in avertisp2p:
        continue

    # Récupération des ref de la machine
    machines = ldap.search('ipHostNumber=%s' % ip, 'w' )['machine']
    if len(machines) == 0:
        # La machine a ete supprimee entre temps
        continue
    machine = machines[0]
    hostname = machine.nom()
    proprio = machine.proprietaire()
    blacklist = proprio.blacklist()

    if isinstance(machine, MachineFixe):
        _, chambre = reperage_chambre(machine.mac()) 
    else:
        chambre = 'machine wifi' 


    # Envoi du mail à disconnect
    mail = connectsmtp()

    if p2p.disconnect_mail:
        requete = "select date from p2p where date > timestamp 'now' - interval '2 hours' and ip_src='%s' order by date limit 1"%ip
        curseur.execute(requete)
        date = curseur.fetchall()[0][0]
        corps = p2p.avertissement % { 'From': upload.expediteur,
                                      'To': upload.expediteur,
                                      'protocole': protocole,
                                      'hostname': hostname,
                                      'nb_paquets': nombre,
                                      'datedebut': date,
                                      'chambre': chambre,
                                      }
        corps = corps.encode('utf-8')
        mail.sendmail(upload.expediteur, upload.expediteur, corps)

    # Inscription dans la base des avertis
    requete = "INSERT INTO avertis_p2p (ip_crans,date,protocole) VALUES ('%s','now','%s')" % (ip, protocole)
    curseur.execute(requete)

    # On envoie un mail a l'adhérent
    ################################
    corps = p2p.deconnexion % { 'From': p2p.expediteur,
                                'To': proprio.email(),
                                'protocole': protocole,
                                'hostname': hostname }
    corps = corps.encode('utf-8')
    mail.sendmail(p2p.expediteur, proprio.email(), corps)

    # L'adhérent n'est pas encore blacklisté
    fin = int(time()) + 24*3600
    proprio.blacklist(['now', fin, 'autodisc_p2p', hostname])
    proprio.save()

    # Vérification du nombre de déconnexions
    #########################################
    nb_decos = len([ x for x in proprio.blacklist() if int(x.split('$')[0]) > time()-365*24*3600 and x.split('$')[2] == 'autodisc_p2p' ])
    if nb_decos >= 3:
        # Génération du fichier postscript
        try:
            fichier_ps = generate_ps('p2p', proprio, ldap)
        except:
            fichier_ps = ("ERREUR lors de la génération. Merci de regénérer manuellement la fiche avec la commande :\n"
                + "/usr/scripts/surveillance/fiche_deconnexion/generate.py --p2p aid=%d" % int(proprio.id()))

        # Envoi du mail à disconnect
        corps = p2p.message_disconnect_multi % { 'from': p2p.expediteur,
                                                 'to': p2p.expediteur,
                                                 'nbdeco': nb_decos,
                                                 'proprio':proprio.Nom(),
                                                 'ps': fichier_ps }
        corps = corps.encode('utf-8')
        mail.sendmail(p2p.expediteur, p2p.expediteur, corps)

################################################################################
#                   Gestion du bittorrent                                      #
################################################################################

# Dans la table p2p on sélectionne les ip_src qui appartiennent au réseau
requete = "SELECT ip_src,count(ip_src) FROM p2p WHERE %s AND date > timestamp 'now' - interval '2 hours' AND id_p2p = 1 GROUP BY ip_src,id_p2p ORDER BY ip_src" % ip_src_in_crans
curseur.execute(requete)
fraudeurs = curseur.fetchall()

# Récupération des fraudeurs pour ne pas les avertir de nouveau
requete = "SELECT ip_crans,protocole FROM avertis_p2p WHERE  date > timestamp 'now' - interval '14 days'"
curseur.execute(requete)
avertisp2p = curseur.fetchall()

for ip, nombre in fraudeurs:
    # On récupére le protocole de p2p :
    protocole = 'Bittorrent'

    # On ne prend pas en compte s'il est sous le seuil admis, ou
    #s'il est averti
    if nombre <= p2p.limite[protocole] or (ip, protocole) in avertisp2p:
        continue

    # Récupération des ref de la machine
    machines = ldap.search('ipHostNumber=%s' % ip, 'w' )['machine']
    if len(machines) == 0:
        # La machine a ete supprimee entre temps
        continue
    machine = machines[0]
    hostname = machine.nom()
    proprio = machine.proprietaire()

    # On envoie un mail a l'adhérent
    ################################
    mail = connectsmtp()

    corps = p2p.avertissement_bt % { 'From': p2p.expediteur,
                                     'To': proprio.email(),
                                     'protocole': protocole,
                                     'hostname': hostname }
    corps = corps.encode('utf-8')
    mail.sendmail(p2p.expediteur, proprio.email(), corps)

    # Inscription dans la base des avertis
    requete = "INSERT INTO avertis_p2p (ip_crans,date,protocole) VALUES ('%s','now','%s')" % (ip, protocole)
    curseur.execute(requete)