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

"""
Script in python for the project in CPL.
Date: 27.12.2019
Last modification: 21.01.2020
Author: Michael da Silva & Nenad Rajic
Version Python: 2.7.15+
Version script: 0.4.3
"""

from scapy.all import *

OWN_ADDR = "0101" #Unicast address of the user
MULTI_ADDR = "1011" #Multicast address of the user
BROADCAST = "1111" #Broadcast address
SIZE_ROW = 7
SIZE_COLUMN = 6
SIZE_ADDR = 4 #Size of destination and source address
MAX_INT_42_BITS = 4398046511103
BITS_FOR_SIZE = 6

"""
Crée un payload contenant des bits aléatoires (juqu'à 42 bits)
"""
def RandomPayload():
    number = random.randint(0, MAX_INT_42_BITS)
    number = "{0:b}".format(number) #Transforme un integer en chiffre binaire
    return number
"""
Corrige le bit érroné grâce au numéro de ligne et numéro de colonne
"""
def CorrectErrorPayload(bits, row, column):
    bits_corrected = bits
    bits_corrected[(row*(SIZE_ROW+1)) + column] = 0 if bits[(row*(SIZE_ROW+1)) + column] else 1
    print("Error corrected")
    return bits_corrected
"""
Ajoute des bits de bourrage au payload s'il n'a pas assez de bits (42 bits de data demandé)
"""
def FillPayload(payload):
    while len(payload) != SIZE_ROW * SIZE_COLUMN:
        payload.append(0)
    return payload
"""
Transforme une liste de int en une chaîne de caractère
"""
def BinaryToString(frame):
    s = ''.join([str(i) for i in frame])
    return s
"""
Transforme une chaîne de caractère en liste de integer
"""
def StringToBinary(frame):
    int_array = list(map(int, frame))
    return int_array
"""
Retourne en binaire la longueur réelle d'un payload (sans bit de bourrage)
"""
def PayloadSize(length):
    binary = "{0:b}".format(length)
    while len(binary) != BITS_FOR_SIZE:
        binary = "0" + binary
    return binary
"""
Vérifie les parités du payload
"""
def HVParityCheck(payload):
    matrix = []
    errors_row = []
    errors_column = []

    while payload: #Création d'une matrice pour faciliter la vérification
        matrix.append(payload[:(SIZE_ROW + 1)])
        payload = payload[(SIZE_ROW + 1):]

    for i in range(len(matrix)):
        if SimpleParityCheck(matrix[i]): #Vérification des parités horizontales
            errors_row.append(i)

    for j in range(len(matrix[0])):
        column = [] #Création d'une colomne pour la vérification des parités verticales
        for i in range(len(matrix)):
            column.append(matrix[i][j])
        if SimpleParityCheck(column):
            errors_column.append(j)

    if len(errors_row) > 1 or len(errors_column) > 1:
        print("Error in data. Can't correct it (2-3 bits error or 8 bits burst)")
        return [-1]
    elif len(errors_row) == 1 & len(errors_column) == 1:
        print("1 error to correct")
        return [errors_row[0], errors_column[0]] #Renvoie des numéros de ligne et colomne
    else:
        print("No error detected")
        return []
"""
Calcule les parités verticales et horizontales pour le payload
"""
def HVParityCalculate(payload):
    result = []
    matrix = []

    for i in range(SIZE_COLUMN):
        matrix.append(SimpleParityCalculate(payload[i*SIZE_ROW:(i+1)*SIZE_ROW])) #Calcul des parités horizontales

    parity_line = []
    parity = 0

    for i in range(SIZE_ROW):
        for row in matrix:
            parity += row[i] #Calcul de la dernière ligne contenant les parités verticales
        parity_line.append(parity % 2)
        parity = 0

    matrix.append(SimpleParityCalculate(parity_line)) #Calcul de la parité de la dernière ligne
    result = [bit for row in matrix for bit in row] #Transformation de la matrice en tableau à une dimension
    return result
"""
Enlève les bits de parité du payload et le renvoie nettoyé
"""
def CleanPayload(payload):
    payload_cleaned = []

    for i in range(SIZE_COLUMN):
        payload_cleaned.append(payload[:(SIZE_ROW)]) #Nouveau payload sans bit de parité
        payload = payload[(SIZE_ROW + 1):]
    result = []
    for i in payload_cleaned: #Nouvelle liste des bits
        for j in i:
            result.append(j)

    return result
"""
Vérifie la parité d'une trame donnée et renvoie si la parité est juste ou non
"""
def SimpleParityCheck(frame):
    parity_to_check = frame[-1]
    parity = 0

    for i in frame[:-1]:
        parity += i

    parity %= 2

    if parity_to_check == parity:
        return 0
    else:
        return -1
"""
Calcule la parité d'une trame donnée et renvoie la trame contenant la parité
"""
def SimpleParityCalculate(frame):
    frame_with_parity = frame

    total = 0
    for i in frame:
        total += i

    frame_with_parity.append(total % 2)
    return frame_with_parity
"""
Envoie automatiquement une payload avec l'adresse de destination, l'adresse source et la longueur du payload (permet d'enlever les bits de bourrage)
"""
def SendFrameAutom(destination, payload):
    length = 0
    destination = StringToBinary(destination)
    source = StringToBinary(OWN_ADDR)
    payload = StringToBinary(payload)

    if len(payload) < SIZE_ROW * SIZE_COLUMN:
        length = len(payload)
        print("Payload too small")
        print("Bits filling...")
        payload = FillPayload(payload) #Bourrage du payload si bits manquant (n'atteint pas la taille de 42 bits)
        print("Done!")
    elif len(payload) == (SIZE_ROW * SIZE_COLUMN):
        length = SIZE_ROW * SIZE_COLUMN

    length = PayloadSize(length) #Taille du payload en binaire
    length = SimpleParityCalculate(StringToBinary(length))

    print("Calculating the parity of destination address...")
    destination = SimpleParityCalculate(destination) #Calcul de la parité de l'adresse de destination
    print("Done!")
    print("Calculating the parity of source address...")
    source = SimpleParityCalculate(source) #Calcul de la parité de l'adresse source
    print("Done!")
    print("Calculating the parity of payload...")
    payload = HVParityCalculate(payload) #Calcul de les parités du payload
    print("Done!")
    #Transformation des informations de binaire à chaîne de caractère (pour la fonction SendFrame)
    destination = BinaryToString(destination)
    source = BinaryToString(source)
    length = BinaryToString(length)
    payload = BinaryToString(payload)

    print("Sending packet...")
    SendFrame(destination, source, length, payload)
"""
Reçoit un paquet et retire les informations telle que l'adresse de destination, l'adresse source, la longueur du payload et le payload
"""
def ReceiveFrameAutom(packet):
    if packet.load[:SIZE_ADDR] == OWN_ADDR or packet.load[:SIZE_ADDR] == MULTI_ADDR or packet.load[:SIZE_ADDR] == BROADCAST: # Vérification de l'adresse de destination
        print("Packet received with one of the valid addresses")
        destination = packet.load[:(SIZE_ADDR+1)]
        source = packet.load[(SIZE_ADDR+1):(SIZE_ADDR+1)*2]
        length = packet.load[(SIZE_ADDR+1)*2:((SIZE_ADDR+1)*2)+(BITS_FOR_SIZE+1)]
        payload = packet.load[((SIZE_ADDR+1)*2)+(BITS_FOR_SIZE+1):]

        print("Checking parity of the destination address...")
        test_dst = SimpleParityCheck(StringToBinary(destination)) # Test de la parité de l'adresse de destination

        if test_dst:
            print("Error in destination, packet dropped")
            return
        else:
            print("Destination correct!")
            destination = destination[:-1]

        test_src = SimpleParityCheck(StringToBinary(source)) # Test de la parité de l'adresse source
        if test_src:
            print("Error in destination, packet dropped")
            return
        else:
            print("Destination correct!")
            source = source[:-1]

        test_size = SimpleParityCheck(StringToBinary(length)) # Test de la parité de la taille du payload
        if test_size:
            print("Error in size of payload, packet dropped")
            return
        else:
            length = length[:-1]

        print("Checking parity of the payload...")
        test_payload = HVParityCheck(StringToBinary(payload)) # Test des parités du payload

        if len(test_payload) == 1: # Rejette le paquet s'il y a plus d'une erreur dans le payload
            return
        elif len(test_payload) == 2:
            payload = CorrectErrorBinary(payload, test_payload[0], test_payload[1])

        print("Cleaning the payload from the parity bits...")
        payload = BinaryToString(CleanPayload(payload))
        print("Done!")

        payload = payload[:int(length, 2)] # Le payload ne garde que les bits d'information reçu, pas les bits de bourrage

        print("\n\n----------Packet received----------\n")
        print("Destination: " + destination)
        print("Source: " + source)
        print("Payload: " + payload)
        print("\n-----------------------------------\n\n")
"""
Envoie une trame via le protocole Ethernet en broadcast (notre protocole gère les adresses de destination)
"""
def SendFrame(destination, source, length, payload):
    data = destination + source + length + payload
    sendp(Ether(dst="ff:ff:ff:ff:ff:ff")/data)
"""
Permet de voir les paquets envoyé sur le réseau et de récupérer ceux ayant comme destinataire l'adresse de broadcast
"""
def Sniff():
    sniff(filter = "ether dst ffffffffffff", prn = ReceiveFrameAutom) # Renvoi du paquet capturé à la fonction ReceiveFrameAutom
"""
Fonction main du script avec quelques tests
"""
def Main():
    """print("Receive(r) or Send(s) ?: ")
    choice = raw_input()

    if choice == "s":
        while 1:
            print("Destination: ")
            dest = raw_input()
            print("Data (if you type enter without data, a random one is created): ")
            data = raw_input()

            if not len(data):
                data = RandomPayload()
            SendFrameAutom(dest, data)
    elif choice == "r":
        Sniff()"""

    # Test pour la détection et correction d'erreur
    """
    data = "00000000000000000000000000000000000000000000000000000000"
    data_1error = "00000000000000000000000000000010000000000000000000000000"
    data_2errors = "00000000000000000000000000000000000010100000000000000000"
    data_3errors = "00000000000000000001011000000000000000000000000000000000"
    data_8burst = "00000000000001000000100000000000000000000000000000000000"

    print(data_8burst)
    test = HVParityCheck(StringToBinary(data_8burst))
    if len(test) == 2:
        payload = CorrectErrorPayload(StringToBinary(data_1error), test[0], test[1])
        print(BinaryToString(payload))
    elif len(test) == 1:
        return
    """
    # Test pour le calcul de parité simple et vérification
    """
    data = "10010"
    test = SimpleParityCalculate(StringToBinary(data))
    print(BinaryToString(test))
    if SimpleParityCheck(test) < 0:
        print("Error in parity")
    else:
        print("Parity correct")
    """
    # Test pour le calcul et la vérification des parités d'un payload
    """
    data = RandomPayload()
    print(data)
    data = FillPayload(data)
    test = HVParityCalculate(StringToBinary(data))
    print(BinaryToString(test))
    HVParityCheck(test)
    """
    # Test pour le calcul de la taille réelle d'un payload (aucun bit de bourrage)
    """
    data = "11"
    length = len(data)
    data = BinaryToString(FillPayload(StringToBinary(data)))
    print(data)
    length = PayloadSize(length)
    data = data[:int(length, 2)]
    print(data)
    """
Main()