from Crypto.Util.number import bytes_to_long, getPrime, isPrime
from Crypto.Random.random import randint
import signal
from secret import FLAG

bit_length = 384
p = getPrime(bit_length)
q = getPrime(bit_length)
n = p**2 * q

while True :
    g = randint(2, n - 1)
    if pow(g, p - 1, p**2) != 1:
        break

h = pow(g, n, n)
r = randint(0, n - 1)

def encrypt(m):
    c = pow(g, m, n) * pow(h, r, n) % n
    return c

L = lambda x : (x - 1) // p
def decrypt(c):
    a = pow(c, p - 1, p**2)
    b = pow(g, p - 1, p**2)

    a = L(a)
    b = L(b)

    m = (a * pow(b, -1, p)) % p
    return m

options = ["Encrypt", "Decrypt", "Get Flag", "Exit"]
def menu() :
    [print(f"{i + 1}. {opt}") for i, opt in enumerate(options)]
    return int(input(">> "))

def main() :
    print(f"{g = }")

    while True :
        choice = menu()
        
        if choice == 1 :
            m = int(input("Enter Your Plaintext: "))
            print(f"Encrypted: {encrypt(m)}")
        elif choice == 2 :
            m = int(input("Enter Your Ciphertext: "))
            print(f"Decrypted: {decrypt(m)}")
        elif choice == 3 :
            challenge = randint(p + 1, p << 10)
            print(f"Challenge: {challenge}")
            response = int(input("Plaintext: "))
            if challenge == response :
                print("Congrats!")
                print(FLAG)
                exit()
            else :
                print("Wrong!")
                exit()
        else :
            print("Kay Bye!")
            exit()

if __name__ == "__main__" :
    signal.alarm(60)
    try :
        main()
    except Exception as e :
        print("Uh Oh")
        exit()