#!/usr/bin/env pypy3

import os
from pydoc import plain
from sys import byteorder
from Crypto.Cipher import AES
from Crypto.Util import Counter
import hashlib

# Create a secret.py file with a variable `FLAG` for local testing :)
from secret import FLAG

secret_key = os.urandom(16)

def encrypt(plaintext, counter):
    m = hashlib.sha256()
    m.update(counter.to_bytes(8, byteorder="big"))

    alg = AES.new(secret_key, AES.MODE_CTR, nonce=m.digest()[0:8])
    ciphertext = alg.encrypt(plaintext)

    return ciphertext.hex()

def xor(bytes1, bytes2):
    return bytes([b1 ^ b2 for b1, b2 in zip(bytes1, bytes2)])

def main():
    print("DES is broken, long live the secure AES encryption!")
    print("Give me a plaintext and I'll encrypt it a few times for you. For more security of course!")

    try:
        plaintext = bytes.fromhex("41" * 512)
        print(plaintext)
    except ValueError:
        print("Please enter a hex string next time.")
        exit(0)
    
    encrypted_plaintext = []

    for i in range(0, 255):
        crypt = encrypt(plaintext, i)
        encrypted_plaintext.append(bytes.fromhex(crypt))
        print(f"Ciphertext {i:03d}: {crypt}")

    keystreams = [xor(enc, plaintext) for enc in encrypted_plaintext]

    enc_flag = encrypt(FLAG.encode("ascii"), int.from_bytes(os.urandom(1), byteorder="big"))
    
    print("Flag:", enc_flag)

    dec_flags = [xor(keystream, bytes.fromhex(enc_flag)) for keystream in keystreams]

    for flag in dec_flags:
        if b'cscg' in flag:
            print(flag)



if __name__ == "__main__":
    main()