#!/usr/bin/env python3

import os
import binascii
import hashlib
import threading
import time
import sys
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

seed_bits = 23
seed_max  = 1 << seed_bits
seed_len  = (seed_bits + 7) // 8
key = os.urandom(16)

def hash_seed(seed_int: int) -> bytes:
    sb = seed_int.to_bytes(seed_len, "big")
    return hashlib.sha256(sb).digest()[:16]

seed  = int.from_bytes(os.urandom(4), "big") % seed_max
seed2 = int.from_bytes(os.urandom(4), "big") % seed_max
K1 = hash_seed(seed)
K2 = hash_seed(seed2)

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

def read_hex(prompt: str):
    s = input(prompt).strip()
    try:
        return binascii.unhexlify(s)
    except Exception:
        print("hmm")
        return None

def enc_cfb(pt: bytes) -> bytes:
    iv = os.urandom(16)
    aes = AES.new(key, AES.MODE_CFB, iv=iv, segment_size=128)
    ct = aes.encrypt(pt)
    return iv + ct

def enc_cbc(data: bytes, iv1: bytes, iv2: bytes, padd: bool) -> bytes:
    x  = pad(data, 16) if padd else data
    c1 = AES.new(K1, AES.MODE_CBC, iv=iv1).encrypt(x)
    c2 = AES.new(K2, AES.MODE_CBC, iv=iv2).encrypt(c1)
    return c2

def menu():
    print("""
1. encrypt 
2. profit 
3. get third 
4. exit
    """)

third = 0
iv11 = None
iv22 = None

def alarm():
    time.sleep(180)
    print("zzz")
    sys.exit(0)

threading.Thread(target=alarm, daemon=True).start()

while True:
    menu()
    op = input("> ").strip()

    if op == "1":
        data = read_hex("pt: ")
        if data is None:
            print()
            continue
        out = enc_cfb(data)
        print("ct: ", out.hex())
        print()

    elif op == "2":
        if iv11 is not None and iv22 is not None:
            iv1, iv2 = iv11, iv22
            iv11 = iv22 = None
        else:
            iv1 = os.urandom(16)
            iv2 = os.urandom(16)
        ct  = enc_cbc(flag, iv1, iv2, padd=True)
        print("iv1: ", iv1.hex())
        print("iv2: ", iv2.hex())
        print("ct: ", ct.hex())
        print()

    elif op == "3":
        if third:
            print("sheesh")
            continue
        block = read_hex("pt: ")
        if block is None:
            print()
            continue
        if len(block) != 16:
            print("hmmm\n")
            continue
        iv1 = os.urandom(16)
        iv2 = os.urandom(16)
        ct  = enc_cbc(block, iv1, iv2, padd=False)
        iv11, iv22 = iv1, iv2
        print("ct: ", ct.hex())
        third = 1
        print()

    elif op == "4":
        break
    else:
        print("mabokkkk?")