import os
from Crypto.Cipher import AES
from Crypto.Util import Counter

def main():
    flag = os.environ.get('FLAG', 'hkcert21{*******************************REDACTED*******************************}')
    flag = flag.encode()
    assert len(flag) == 80

    key = os.urandom(16)
    iv = os.urandom(16)

    options = ['ecb', 'cbc', 'cfb', 'ofb', 'ctr']
    suboptions = ['data', 'flag']
    
    for _ in range(5):
        [option, suboption, *more] = input('> ').split(' ')
        if option not in options: raise Exception('invalid option!')
        if suboption not in suboptions: raise Exception('invalid suboption!')
        options.remove(option)

        if suboption == 'data':
            message = bytes.fromhex(more[0])
        else:
            message = flag
        
        if option == 'ecb':   cipher = AES.new(key, AES.MODE_ECB)
        elif option == 'cbc': cipher = AES.new(key, AES.MODE_CBC, iv)
        elif option == 'cfb': cipher = AES.new(key, AES.MODE_CFB, iv, segment_size=128)
        elif option == 'ofb': cipher = AES.new(key, AES.MODE_OFB, iv)
        elif option == 'ctr': cipher = AES.new(key, AES.MODE_CTR, counter=Counter.new(16, prefix=iv[:14]))

        ciphertext = cipher.encrypt(message)
        print(ciphertext.hex())
    else:
        print('Bye!')

if __name__ == '__main__':
    main()
