#!/usr/bin/env python3

import os
from Pailier import *
from Crypto.Util.number import *

with open("/flag.txt", "rb") as f:
    flag_bytes = f.read()

key = os.urandom(66)              
key_int = bytes_to_long(key)

cipher = pailier()

while True:
    print("1. encrypt")
    print("2. bingo")
    print("3. decrypt")
    print("4. key?")
    try:
        inp = int(input("> "))
    except (ValueError, EOFError):
        print("Invalid input")
        continue

    if inp == 1:
        print("pt (hex)")
        try:
            inp = input("> ")
            ct = cipher.encrypt(int(inp, 16))
            print('ct : ', '{0:x}'.format(ct))
        except (ValueError, EOFError):
            print("Invalid hex input")

    elif inp == 2:
        print("key (hex)")
        try:
            user_hex = input("> ").strip()
            user_key = bytes.fromhex(user_hex)
            
            if len(user_key) == 66 and user_key == key:
                try:
                    print(flag_bytes.decode())
                except Exception:
                    print(flag_bytes.hex())
            else:
                print("nope")
        except (ValueError, EOFError):
            print("nope")

    elif inp == 3:
        print("ct (hex)")
        try:
            inp = input("> ")
            pt = cipher.decrypt(int(inp, 16))
            print('pt : ', '{0:x}'.format(pt))
        except (ValueError, EOFError):
            print("Invalid hex input")

    elif inp == 4:
        try:
            ct = cipher.encrypt(key_int)
            print('ct : ', '{0:x}'.format(ct))
        except Exception:
            print("Encryption error")

    else:
        exit()