# referencia: https://www.geeksforgeeks.org/implementation-diffie-hellman-algorithm/

# Código de Diffie-Hellman

def verificador_primo(p):
    # Comprueba si el número ingresado es un número primo o no
    if p < 1:
        return -1
    elif p > 1:
        if p == 2:
            return 1
        for i in range(2, p):
            if p % i == 0:
                return -1
        return 1

def verificador_primitivo(g, p, L):
    # Comprueba si el número ingresado es una raíz primitiva o no
    for i in range(1, p):
        L.append(pow(g, i) % p)
    for i in range(1, p):
        if L.count(i) > 1:
            L.clear()
            return -1
    return 1

l = []
while 1:
    P = int(input("Ingrese P: "))
    if verificador_primo(P) == -1:
        print("El número no es primo. ¡Por favor, inténtelo de nuevo!")
        continue
    break

while 1:
    G = int(input(f"Ingrese la raíz primitiva de {P}: "))
    if verificador_primitivo(G, P, l) == -1:
        print(f"{G} no es una raíz primitiva de {P}. ¡Por favor, inténtelo de nuevo!")
        continue
    break

# Claves privadas
x1, x2 = int(input("Ingrese la clave privada del Usuario 1: ")), int(
    input("Ingrese la clave privada del Usuario 2: "))
while 1:
    if x1 >= P or x2 >= P:
        print(f"Ambas claves privadas deben ser menores que {P}!")
        continue
    break

# Calcular las claves públicas
y1, y2 = pow(G, x1) % P, pow(G, x2) % P

# Generar las claves secretas
k1, k2 = pow(y2, x1) % P, pow(y1, x2) % P

print(f"\nLa clave secreta para el Usuario 1 es: {k1}\nLa clave secreta para el Usuario 2 es: {k2}\n")

if k1 == k2:
    print("Las claves han sido intercambiadas correctamente.")
else:
    print("Las claves no han sido intercambiadas correctamente.")

# Ejemplos de valores que se pueden probar
# P = 23, raíz primitiva G = 5:

# Clave privada del Usuario 1 (x1) = 7
# Clave privada del Usuario 2 (x2) = 8


#Para P = 17, raíz primitiva G = 3:

# Clave privada del Usuario 1 (x1) = 10
# Clave privada del Usuario 2 (x2) = 12

# Para P = 29, raíz primitiva G = 2:

# Clave privada del Usuario 1 (x1) = 15
# Clave privada del Usuario 2 (x2) = 20