#!/usr/bin/env python3
# local_exploit.py
import json,os
from Crypto.Cipher import AES

### Global variables
secret_key = os.urandom(16)
ctr_nonce = os.urandom(8)

def gen_user_cookie():
    cookie_dict = {}
    cookie_dict["access_level"] = "User"
    pt = json.dumps(cookie_dict).encode()
    cipher = AES.new(secret_key, AES.MODE_CTR, nonce = ctr_nonce)
    ct = cipher.encrypt(pt) # Made it return bytes instead of hex encoded
    return pt, ct # Made it return the plaintext too

def decrypt_cookie(cookie):
    ct = bytes.fromhex(cookie)
    cipher = AES.new(secret_key, AES.MODE_CTR, nonce = ctr_nonce)
    pt = cipher.decrypt(ct)
    return pt

# Make cookie and plaintext
plaintext, cookie = gen_user_cookie()
print(f"Plaintext (bytes): {plaintext}")
print(f"Plaintext (hex): {plaintext.hex()}")
print(f"Encrypted Cookie (hex): {cookie.hex()}")

# Calculate the keystream used in xor
keystream = []

print("\nCalculating keystream...")
for pt, ct in zip(plaintext, cookie):
    s = pt ^ ct
    keystream.append(s)
    print(f"{pt:02x} ^ {ct:02x} = {s:02x}")

# Encrypt new plaintext
new_dict = {}
new_dict["access_level"] = "Admin"
new_plaintext = json.dumps(new_dict).encode()

print(f"\nNew plaintext (bytes): {new_plaintext}")
print(f"\nNew plaintext (hex): {new_plaintext.hex()}")

new_ciphertext = ""

print("\nEncrypting new plaintext...")
for s, pt in zip(keystream, new_plaintext):
    ct = s ^ pt
    new_ciphertext += f"{ct:02x}"
    print(f"{s:02x} ^ {pt:02x} = {ct:02x}")

print(f"New Ciphertext (hex): {new_ciphertext}")

print(f"\nDecrypted Ciphertext with known secret_key and nonce: {decrypt_cookie(new_ciphertext)}")

# Bruteforce the last (extra) byte
print("\nBruteforcing the last (extra) byte")
for i in range(2**8):
    ct = new_ciphertext + f"{i:02x}"
    decrypted = decrypt_cookie(ct)
    print(f"{ct} decrypted to {decrypted}", end="...")
    if decrypted == new_plaintext:
        print("Correct")
        new_ciphertext = ct
        break
    print("Wrong")

print(f"\nNew Ciphertext with correct extra byte (hex): {new_ciphertext}")