from time import time
import random

def get_num_random():
    return random.randrange(2, 99999)

def exponencial(a, b):
    if b == 0:
        return 1
    else:
        res = exponencial(a, int(b/2))
        if (b % 2) == 0:
            return res*res
        else:
            return res*res*a

def lagrange(a, n, p):
    a = a % p
    n = n % (p-1)
    if n == 0:
        return (1 % p)
    elif n >= 1:
        x = 1
        aux_x = a
        while n != 0:
            delta = n % 2
            x = (x * exponencial(aux_x, delta))
            x = x % p
            n = (n - delta) / 2
            aux_x = exponencial(aux_x, 2)
        return x

def modulo_A(p, g):
    Kprivada_A = get_num_random()  # Llave privada
    Kpublica_A = lagrange(g, Kprivada_A, p)  # Lave publica
    return Kprivada_A, Kpublica_A

def modulo_B(p, g):
    Kprivada_B = get_num_random()  # Llave privada
    Kpublica_B = lagrange(g, Kprivada_B, p)  # Lave publica
    return Kprivada_B, Kpublica_B

def encriptar_msje(Kprivada, Kpublica, p):
    mensaje = int(input("\tMensaje : "))
    param = lagrange(Kpublica, Kprivada, p)
    mensaje = mensaje * param
    mensaje = mensaje % p
    print("\tMensaje encriptado : ", mensaje)
    return mensaje

def desencriptar_msje(mensaje, kpublica, kprivada, p):
    exp = (p - 1) - kprivada
    param = lagrange(kpublica, exp, p)
    mensaje = mensaje * param
    mensaje = mensaje % p
    print("\tMensaje desencriptado : ", mensaje)
    
# ----------------- main ---------------------------
print("\t\nAlgoritmo ElGamal - Sistema de Encriptacion\n")
p = int(input("\tNumero Primo [ p ] : "))
g = int(input("\tGenerador    [ g ] : "))
Kprivada_A, Kpublica_A = modulo_A(p, g)
Kprivada_B, Kpublica_B = modulo_B(p, g)
Termina = "n"

while Termina != "y":
    print("\n\t1) Mandar mensaje de A --> B")
    print("\t2) Mandar mensaje de B --> A\n")

    opcion = int(input("\tOpcion [1,2] : "))
    print()
    if opcion == 1:
        inicio_enc = time()
        mensaje = encriptar_msje(Kprivada_A, Kpublica_B, p)
        final_enc = time()
        print("\tLlave Privada (A) : ",Kprivada_A,"\n\tLlave Publica (A) : ",Kpublica_A)
        print(f"\tTiempo de encriptacion :",(final_enc-inicio_enc),"segundos")
        desencriptar = input("\n\tDesencriptar mensaje [y / n]: ")
        print()
        if desencriptar == "y":
            inicio_des = time()
            desencriptar_msje(mensaje, Kpublica_A, Kprivada_B, p)
            final_des = time()
            print(f"\tTiempo de desencriptacion :",(final_des-inicio_des),"segundos")
            print()
        elif desencriptar == "n":
            print("\t\nINFO : Mensaje,Kprivada_A,Kpublica_A,Kprivada_B,Kpublica_B,Primo,Generador")
            print("\t",mensaje, Kprivada_A, Kpublica_A, Kprivada_B, Kpublica_B,p,g)
    elif opcion == 2:
        inicio_enc = time()
        mensaje = encriptar_msje(Kprivada_B, Kpublica_A, p)
        final_enc = time()
        print("\tLlave Privada (B) : ",Kprivada_B,"\n\tLlave Publica (B) : ",Kpublica_B)
        print(f"\tTiempo de encriptacion :",(final_enc-inicio_enc),"segundos")
        desencriptar = input("\n\tDesencriptar mensaje [y / n]: ")
        print()
        if desencriptar == "y":
            inicio_des = time()
            desencriptar_msje(mensaje, Kpublica_B, Kprivada_A, p)
            final_des = time()
            print(f"\tTiempo de desencriptacion :",(final_des-inicio_des),"segundos")
            print()
        elif desencriptar == "n":
            print("\t\nINFO: Mensaje,Kprivada_A,Kpublica_A,Kprivada_B,Kpublica_B,Primo,Generador")
            print("\t",mensaje, Kprivada_A, Kpublica_A, Kprivada_B, Kpublica_B, p,g)
    Termina = input("\n\tTerminar programa [y / n]: ")
    print()