from samson.utilities.bytes import Bytes
from samson.oracles.default_oracle import DefaultOracle

class CBCIVKeyEquivalenceAttack(object):
    """
    Performs a key-recovery attack.

    Some systems may opt to reduce the number of parameters to decrypt a ciphertext by
    setting the IV so it matches the key. Even assuming they don't include the IV with
    the ciphertext, we can recover the key under certain conditions.

    Conditions:
        * CBC is being used
        * The IV and the key are equivalent
        * The user has access to an oracle that returns the decrypted plaintext
    """
    def __init__(self, oracle: DefaultOracle, block_size: int=16):
        """
        Parameters:
            oracle  (DefaultOracle): An oracle that takes in a bytes-like object and returns its decryption.
            block_size        (int): The block size of the block cipher in bytes.
        """
        self.oracle = oracle
        self.block_size = block_size


    def execute(self, ciphertext: bytes) -> Bytes:
        """
        Executes the attack.
        
        Parameters:
            ciphertext (bytes): Bytes-like ciphertext generated by the system.

        Returns:
            Bytes: The key to the cipher.
        """
        edited_cipher = Bytes.wrap(bytes(ciphertext))

        edited_cipher[self.block_size:self.block_size * 2] = b'\x00' * self.block_size
        edited_cipher[self.block_size * 2:self.block_size * 3] = edited_cipher[:self.block_size]

        recovered_bytes = self.oracle.request(bytes(edited_cipher))
        blocks = Bytes(recovered_bytes).chunk(self.block_size)

        derived_iv = blocks[0] ^ blocks[2]
        return derived_iv