from hades import HADES
from Crypto.Util.number import long_to_bytes, bytes_to_long
import os
import signal

with open('flag.txt', 'r') as f :
    FLAG = f.read()

key1 = os.urandom(16)
key2 = os.urandom(16)
cipher = HADES(key1, key2)

def fail() :
    print("Fail!")
    exit()

def menu() :
    print("\n".join(["1. Encrypt", "2. Decrypt", "3. Exit"]))
    resp = input(">> ")
    return int(resp)

def encrypt() :
    pt = bytes.fromhex(input("Enter your plaintext: ").strip())

    if (len(pt) != 16) :
        fail()
    elif (bytes_to_long(pt).bit_count() > 1) :
        fail()
    
    print(f"Encrypted: {cipher.encrypt_block(pt).hex()}")

def decrypt() :
    ct = bytes.fromhex(input("Enter your plaintext: ").strip())

    if (len(ct) != 16) :
        fail()

    pt = cipher.decrypt_block(ct)

    if pt == b"Hwelp Pweaseeeee" :
        print(f"Sure thing comrade: {FLAG}")
        exit()
    else :
        print(f"Me don't understand what u mean")

def main() :
    while True :
        choice = menu()

        if choice == 1 :
            encrypt()
        elif choice == 2 :
            decrypt()
        elif choice == 3 :
            exit()
        else :
            fail()

if __name__ == "__main__" :
    signal.alarm(120)
    try:
        main()
    except Exception as e :
        print("See u Lator")