from samson.utilities.bytes import Bytes
from samson.oracles.encryption_oracle import EncryptionOracle

class OCBAuthForgeryAttack(object):
    """
    Performs a authentication tag forgery attack.

    Conditions:
        * OCB2 is being used
        * The user has access to an oracle that accepts arbitrary plaintext and returns the ciphertext
    """

    def __init__(self, oracle: EncryptionOracle, block_size: int=16):
        """
        Parameters:
            oracle (EncryptionOracle): 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, plaintext: bytes) -> (Bytes, Bytes):
        """
        Executes the attack.
        
        Parameters:
            plaintext (bytes): Bytes-like ciphertext generated by the system.

        Returns:
            (Bytes, Bytes): The key to the cipher.
        """
        encoded_length = Bytes(128).zfill(16)
        _tag, ciphertext = self.oracle.encrypt(encoded_length + plaintext, b'')

        new_ciphertext = ciphertext[:16] ^ encoded_length
        new_tag = plaintext ^ ciphertext[16:]

        return new_tag, new_ciphertext
