test.py

from Crypto.Cipher import AES
from Crypto import Random
from Crypto.Util.Padding import pad, unpad
from typing import Tuple, Optional

from solve import PaddingOracleSolver

BLOCK_SIZE = AES.block_size

key = Random.get_random_bytes(BLOCK_SIZE)

def encrypt(msg: bytes) -> Tuple[bytes, bytes]:
    """
    Encrypts a given message using AES in CBC mode with random IV.
    
    Args:
        msg (bytes): The plaintext message to encrypt.
        
    Returns:
        Tuple[bytes, bytes]: A tuple containing the IV and the ciphertext.
    """
    iv = Random.get_random_bytes(BLOCK_SIZE)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    ct = cipher.encrypt(pad(msg, BLOCK_SIZE))
    return iv, ct

def decrypt(iv: bytes, ct: bytes) -> Optional[bytes]:
    """
    Decrypts a ciphertext using AES in CBC mode and verifies the padding.
    
    Args:
        iv (bytes): The initialization vector used during encryption.
        ct (bytes): The ciphertext to decrypt.
        
    Returns:
        Optional[bytes]: The decrypted plaintext message, or None if padding is invalid.
    """
    cipher = AES.new(key, AES.MODE_CBC, iv)
    pt = cipher.decrypt(ct)
    try:
        return unpad(pt, BLOCK_SIZE)
    except ValueError:
        return None

def oracle(iv: bytes, ct: bytes) -> bool:
    """
    Padding Oracle function that checks if a ciphertext can be decrypted correctly.
    
    Args:
        iv (bytes): The initialization vector used during encryption.
        ct (bytes): The ciphertext to check.
        
    Returns:
        bool: True if the ciphertext can be decrypted correctly (valid padding), False otherwise.
    """
    return decrypt(iv, ct) is not None


solver = PaddingOracleSolver(oracle)

message = Random.get_random_bytes(32)
print(f'message: {message.hex()}')

iv, ciphertext = encrypt(message)
print(f'iv: {iv.hex()}')
print(f'ciphertext: {ciphertext.hex()}')

plaintext = unpad(solver.solve(iv, ciphertext), BLOCK_SIZE)
print(f'plaintext: {plaintext.hex()}')

assert plaintext == message
print('[+] Success')