""" --------------------------------------------------------------------------------------
   Programa que implementa o servidor da comunicacao TCP/IP com Diffie-Hellman
   Objetivo: Comunicacao de cliente-servidor fazendo o estabelecimento de Chave Secreta
   com Diffie-Hellman
   Restricoes: o programa necessita que a porta 10000 nao esteja sendo utilizada.

   Autor: Brendon e Marllon.
   Disciplina: Redes II
   Data da ultima atualizacao: 28/07/2021
----------------------------------------------------------------------------------------"""
import socket
from common.util import *
from des import DesKey


def listen(host, port):
    """ Create a TCP/IP socket listen for incoming connections """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    # Bind the socket to the port
    print('[SERVER LOG] starting up on {} port {}'.format(host, port))
    server_address = (host, port)
    sock.bind(server_address)

    sock.listen(1)
    return sock

def generate_prime_generator(prime_module):
    """ Generate a random prime number to be used as generator number (g) for Diffie-Hellman """
    g = generate_random_prime()
    while g == prime_module:
        g = generate_random_prime()
    
    print('[SERVER LOG] generate prime generator (g) with the value equal: {}'.format(g))
    return g


def send_message(conn, message) -> None:
    """ Send message to client """
    print('[SERVER LOG] sending message to client: {}'.format(str(message)))
    if type(message) == bytes:    
        conn.sendall(message)
    else:
        conn.sendall(str.encode(str(message)))

def generate_random_private_key() -> int:
    """ Generate a random private key """
    private_key = random.randint(1, 10000)
    print('[SERVER LOG] generate random private key equal: {}'.format(str(private_key)))
    return private_key


def calculate_server_result(base, exp, mod) -> int:
    """ Calculate server result """
    result = power(base, exp, mod)
    print('[SERVER LOG] calculated public result to send for the client: {}'.format(result))
    return result


def calculate_private_shared_key(client_result, server_private_key, p) -> int:
    """ Calculate shared key """
    shared_key = power(client_result, server_private_key, p)
    print('[SERVER LOG] calculated shared key equal: {}'.format(shared_key))
    return shared_key


def close_connection(conn):
    """ Clean up the connection """
    print('[SERVER LOG] closing connection')
    conn.close()    


def main():
    sock = listen(host='localhost', port=10000)

    while True:
        # Wait for a connection
        print('[SERVER LOG] waiting for a connection')
        connection, client_address = sock.accept()

        try:
            print('[SERVER LOG] connection from' + str(client_address))

            # Defines Diffie-Hellman shared parameters communicating with connected server
            print('[SERVER LOG] defining Diffie-Hellman parameters with server')

            p = int(connection.recv(1024).decode(encoding='UTF-8')) # Random prime generated by client to be used as module number (p) for Diffie-Hellman
            print('[SERVER LOG] received generate prime module (p) with the value equal: {}'.format(p))

            g = generate_prime_generator(p) # Randomly generated by server to be used as generator (g) for Diffie-Hellman
            send_message(connection, g)

            private_key = generate_random_private_key()

            client_result = int(connection.recv(1024).decode(encoding='UTF-8'))
            print('[SERVER LOG] received server_result with value equal: {}'.format(client_result))

            result = calculate_server_result(g, private_key, p)
            send_message(connection, result)

            shared_private_key = calculate_private_shared_key(client_result, private_key, p)

            des_key = DesKey(shared_private_key.to_bytes(8, byteorder='big'))

            # Receive the data in chunks and retransmit it
            while True:
                data = connection.recv(1024)
                print('[SERVER LOG] received message from the client: {}'.format(data))
                print('[SERVER LOG] received message from the client decrypted: {}'.format(des_key.decrypt(data, padding=True)))
                if len(data) > 0:
                    print('[SERVER LOG] sending data back to the client')
                    send_message(connection, data)
                else:
                    print('[SERVER LOG] no more data from {}'.format(client_address))
                    break
                
        finally:
            close_connection(connection)


if __name__ == "__main__":
    print("===========================================================================")
    print("Inicio da execucao: programa que implementa o server da comunicacao TCP/IP.")
    print("Prof. Elias P. Duarte Jr.  -  Disciplina Redes II")
    print("Autores: Brendon e Marllon")
    print("===========================================================================")
    main()