import zlib
import random
import string
from Crypto.Cipher import ARC4

"""
    cipher = RC4(plaintext)
    KEY is random
    there is no handshake (no need)
"""

# compress and then use RC4 stream cipher to encrypt the message
def encrypt(msg):
    data = msg
    cipher = ARC4.new(KEY)
    compressed_msg = zlib.compress(data)
    return cipher.encrypt(compressed_msg)

# we use the threashold value to prevent infinite recursion (cookie length <= 20 characters)
def find_cookie(curr_cookie, threshold=20):

    found_character = False

    # we try all possible characters, they have values within [33, 127]
    for i in range(33,127):
        # the character we are inserting
        curr_char = chr(i)

        # we use random string to simulate the messages before and after the cookie
        msg_before_cookie = ''.join(random.sample(string.ascii_lowercase + string.digits, k=17))
        msg_after_cookie = ''.join(random.sample(string.ascii_lowercase + string.digits, k=17))

        # create two payloads and encrypt them using RC4 stream cipher
        # '~#:/[|/ç' is a random string that does not appear in the secret
        # TODO: fix the code
        payload1 = msg_before_cookie + "cookie=" + ??? + ??? + '~#:/[|/ç' + ' ' + SECRET.decode() + ' ' + msg_after_cookie
        payload2 = msg_before_cookie + "cookie=" + ??? + '~#:/[|/ç' + ??? + ' ' + SECRET.decode() + ' ' + msg_after_cookie
        enc1 = encrypt(payload1.encode())
        enc2 = encrypt(payload2.encode())

        # TODO: compare the encrypt messages size, and recursively call find_cookie here
        # Hint: if your guess is correct, the newly inserted character will be compressed


    if len(curr_cookie) > threshold or not found_character:
        print('[+] cookie=' + curr_cookie)

if __name__ == '__main__':
    print("{-} CRIME Proof of Concept Attack\n")
    KEY = ''.join(random.sample(string.ascii_uppercase + string.digits, k=17)).encode()
    SECRET = b"cookie=ilovecsc427andarnold"
    print("[+] Secret TOKEN :", SECRET.decode())
    print("[+] Encrypted with \033[33mRC4\033[0m")
    print("[+] Trying to decrypt with a compression oracle attacks using a \033[33mrecursive two_tries\033[0m method\n")
    print("[+] Possible cookies:")
    find_cookie("")