solve.py

from Crypto.Cipher import AES
from typing import Callable

class PaddingOracleSolver:
    """
    Padding Oracle Solver class.
    
    The solver attempts to decrypt a ciphertext encrypted using AES in CBC mode
    by exploiting the padding oracle vulnerability.
    """
    
    def __init__(self, oracle: Callable[[bytes, bytes], bool], block_size: int = AES.block_size):
        """
        Initializes the PaddingOracleSolver instance.
        
        Args:
            oracle (Callable[[bytes, bytes], bool]): The oracle function to query.
            block_size (int): The block size used in AES encryption (default is 16 bytes for AES).
        """
        self.oracle = oracle
        self.block_size = block_size

    def solve(self, iv: bytes, ct: bytes) -> bytes:
        """
        Solves the Padding Oracle attack for the given IV and ciphertext.
        
        Args:
            iv (bytes): The initialization vector used during encryption.
            ct (bytes): The ciphertext to decrypt.
        
        Returns:
            bytes: The decrypted plaintext message (without padding).
        """
        # Combine IV and ciphertext into a single message
        msg = iv + ct
        no_of_blocks = len(msg) // self.block_size
        # Split the message into blocks
        blocks = [msg[i * self.block_size : (i+1) * self.block_size] for i in range(no_of_blocks)]
        plaintext = b''
        
        # Iterate over each block except the last one to perform the attack
        for i in range(len(blocks) - 1):
            leaked = [0] * self.block_size  # Stores the decrypted bytes of the current block
            fakeblock = [0] * self.block_size  # Temporary block for testing guesses

            # Try each byte in the current block and use the oracle to check correctness
            for j in range(1, self.block_size + 1):
                for k in range(0x100):
                    fakeblock[-j] = k
                    if self.oracle(bytes(fakeblock), blocks[i+1]):
                        leaked[-j] = j ^ blocks[i][-j] ^ k  # Recover the plaintext byte
                
                # Adjust the fakeblock to prepare for next byte leak
                for l in range(1, j + 1):
                    fakeblock[-l] = j + 1 ^ blocks[i][-l] ^ leaked[-l]

            plaintext += bytes(leaked)

        # Return the raw plaintext (padded)
        return plaintext