from utils.digit_conversion import *
from utils.dijkstra import *
from Host import Host
from Key import Key
import socket
import configparser


class Sender(Host):
    """docstring for Sender."""
    def __init__(self, config_file):
        super(Sender, self).__init__(config_file)
        self.route = []
        self.KeyID_key = {}
        self.IP_KeyID = {}
        self.g = 2
        self.p = 179769313486231590770839156793787453197860296048756011706444423684197180216158519368947833795864925541502180565485980503646440548199239100050792877003355816639229553136239076508735759914822574862575007425302077447712589550957937778424442426617334727629299387668709205606050270810842907692932019128194467627007
        self.init_keys_done = 0
        self.listen()
        self.write()

    def init_encryption_parameters(self, config_file):
        config = configparser.ConfigParser()
        config.read(config_file)
        parameters = config['parameters']
        self.g = int(parameters['g'])
        self.p = int(parameters['p'])

    def shortest_path(self, topology, destination):
        '''Use Dijkstra's algorithm to find the shortest path between Sender
            and destination given a graph of class Topology : <topology>
            <destination> is a tuple (ip : str, port : int)'''
        if destination not in topology.nodes:
            print('Shortest path research failed:')
            print('{} not in network'.format(destination))
        else:
            weigths, previous_path = dijkstra(topology, (self.ip_addr,self.port_in))
            path = []
            node = destination
            while node is not None:
                path.append(node)
                node = previous_path[node]
            path.reverse()
            return path

    def initialyze_keys(self,path):
        '''Initialyze all the keys with the relays and the receiver listed in the path
            path should correspond to the shortest path found using shortest_path method'''
        for i in range(1,len(path)):
            self.generate_and_send_new_key(path[i][0],path[i][1])

    def generate_and_send_new_key(self, ip_address, port):
        '''1) Generate a Key object in which will be stored a unique Key ID and a public key generated by the Sender.
           2) Send the key ID and the public key to the other entity'''
        key_id = self.generate_key_id()
        new_key = Key(self.g, self.p) # generate a key with the specified g and p parameters for the diffie-hellman algorithm
        self.KeyID_key.update({key_id:new_key}) # update the dico with the link: Key ID - Object Key
        self.IP_KeyID.update({ip_address:key_id}) # update the dico with the link: IP Address - Key ID
        self.send_key_init(ip_address, port, key_id, new_key.get_public_key()) # Send a KEY_INIT message

    def generate_key_id(self):
        '''Generate a unique 32-bits key ID.'''
        key_id = len(self.KeyID_key)+1
        return dec_to_32bits(key_id)

    def send_key_init(self, ip_address, port, key_id, public_key):
        '''Send a KEY_INIT message in the specified format for this project.'''
        version = '0001'
        msg_type = '0000'
        msg_empty_space = '00000000'

        body = key_id + dec_to_4bits(self.g) + dec_to_1024bits(self.p) + dec_to_1024bits(public_key)

        msg_length, body_with_padding = self.compute_msg_length(body)

        header = version + msg_type + msg_empty_space + msg_length
        self.handshake(header + body_with_padding, ip_address,port)

    def handshake(self, msg, ip, port):
        '''Initialyze the key with a specified relay or receiver'''
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((self.ip_addr, self.port_out))
        s.connect((ip, port))
        s.send(msg.encode('utf-8'))

        BUFFER_SIZE = 4096

        data = s.recv(BUFFER_SIZE)
        self.on_data(data, s)
        s.close()

    def send_shallot(self, ip_address, port, shallot):
        '''Send the shallot in the specified format for this project.'''
        version = '0001'
        msg_type = '0010'
        msg_empty_space = '00000000'

        msg_length, shallot_with_padding = self.compute_msg_length(shallot)

        header = version + msg_type + msg_empty_space + msg_length
        self.send(header + shallot_with_padding, ip_address,port)

    def generate_key_from_replier(self, conn, message):
        '''Generate the shared key between the sender and the replier.'''
        key_ID = message[32:64]
        if self.check_key_ID_exist(key_ID):
            self.KeyID_key[key_ID].generate_shared_key(int(message[64:1088],2))
            self.init_keys_done += 1
        else:
            print('ERROR')
            print('INVALID_KEY_ID')
            self.send_error(conn, 1)

    def encrypt(self, key_id, message):
        '''Encrypt the message using the AES algorithm'''
        return self.KeyID_key[key_id].cipher.encrypt(message)

    def decrypt(self, key_id, message):
        '''Decrypt the message using the AES algorithm'''
        return self.KeyID_key[key_id].cipher.decrypt(message)

    def build_shallot(self, path, message):
        '''Build the shallot based on the path'''
        shallot = message
        path.reverse()

        # Build the first layer of the sallot for the Receiver
        IP = path[0][0] # IP used to find the Key ID in the dico self.IP_KeyID
        PORT = path[0][1]
        key_ID = self.IP_KeyID[IP] # Key ID used to find the shared key with the relay/receiver in the dico KeyID_key = {}
        IP_next = path[0][0] # IP_next(Receiver) = IP(Receiver) when the receiver will decipher the shallot
        PORT_next = path[0][1]
        binary_IP_next = ip2bin(IP_next)
        binary_PORT_next = dec_to_32bits(PORT_next)
        shallot = self.encrypt(key_ID,binary_IP_next+binary_PORT_next+shallot)
        shallot = key_ID + shallot

        # Build the following layer for the relays
        for i in range(1,len(path)-1):
            IP = path[i][0] # IP of the current relay
            PORT = path[i][1]
            key_ID = self.IP_KeyID[IP]

            IP_next = path[i-1][0] # IP of the next relay or the receiver)
            PORT_next = path[i-1][1]
            binary_IP_next = ip2bin(IP_next)
            binary_PORT_next = dec_to_32bits(PORT_next)
            shallot = self.encrypt(key_ID,binary_IP_next+binary_PORT_next+shallot)

            shallot = key_ID + shallot

        path.reverse()
        return shallot

    def check_key_ID_exist(self, key_ID):
        '''Check if the Key ID is valid'''
        return key_ID in self.KeyID_key

    def on_data(self, data, conn):
        data = str(data)[2:]
        version=data[0:4]
        msg_type=data[4:8]
        msg_length=data[16:32]
        if msg_type == '0001':
            # MSG TYPE = KEY_REPLY
            print('KEY_REPLY')
            self.generate_key_from_replier(conn, data)
        elif msg_type == '0011':
            # MSG TYPE = ERROR
            print('ERROR')
            if int(data[32:48],2) == 0:
                print('INVALID_MESSAGE_FORMAT')
            elif int(data[32:48],2) == 1:
                print('INVALID_KEY_ID')
        else:
            print('ERROR')
            print('INVALID_MESSAGE_FORMAT')
            self.send_error(conn, 0)
            