from __future__ import division
import math
import random
from Crypto.Cipher import AES

# helpers

def xor(arr1, arr2):
    """XOR two byte arrays. Shorter array should be passed second."""
    if len(arr2) < len(arr1):
        l1, l2 = len(arr1), len(arr2)
        arr2 = arr2 * int(l1 / l2) + arr2[:l1 % l2]
    return bytearray(c1 ^ c2 for c1, c2 in zip(arr1, arr2))

def gen_rand_key(size=16):
    """Generate random key (default 16 bytes)."""
    return bytearray([random.SystemRandom().randint(0, 255)
                      for _ in xrange(size)])

# padding

def pad_pkcs7(key, block_len=16):
    """Pad given key using PKCS#7.
    Args
        key: byte array of key to pad
        block_len: length of block to which to pad key
    Returns
        bytearray of padded key
    """
    if block_len == len(key):
        pad_len = block_len
    else:
        pad_len = block_len - len(key)
    return key + bytearray([pad_len] * pad_len)

def strip_pkcs7(input, raise_error=True):
    """Strip PKCS7 padding from input. Raise error if invalid padding.
    Call with raise_error=False for testing with assertions.
    """
    padding = input[-1]
    valid = True
    for i in xrange(1, padding + 1):
        if i > len(input):
            if raise_error:
                raise IndexError('Padding too long for input')
            valid = False
            break
        if input[-i] != padding:
            if raise_error:
                raise ValueError('Invalid padding detected')
            valid = False

    return input[:-padding] if valid else False

def gen_blocks(code, pad=pad_pkcs7, size=16):
    """Generate blocks from given bytearray code and pad function."""
    num_blocks = int(math.ceil(len(code) / size))
    blocks = [code[i * size: i * size + size] for i in xrange(num_blocks)]

    # PKCS#7 adds extra block if last code block matches block size
    if pad == pad_pkcs7:
        padded = pad(blocks[-1], size)
        last = ([padded] if size != len(blocks[-1]) else
                [padded[:size], padded[size:]])
    # no padding specified
    else:
        last = [blocks[-1]]

    return blocks[:-1] + last

# ECB

def apply_ECB(mode, input, key, pad=pad_pkcs7):
    """Apply ECB to given key and inputs.
    Args
        mode: str, 'encrypt' or 'decrypt'
        input: bytearray of input code
        key: AES key object in ECB mode
        pad: padding function to use on input
    Returns
        Bytearray of encrypted or decrypted input.
    """
    blocks = gen_blocks(input, pad if mode == 'encrypt' else None)
    aes_f = key.encrypt if mode == 'encrypt' else key.decrypt
    input_str = ''.join([str(block) for block in blocks])
    return bytearray(aes_f(input_str))

# CBC

def CBC_encrypt(aes_f, blocks):
    """CBC encrypt: first XOR then encrypt."""
    prev = blocks[0]
    output = bytearray()
    for block in blocks[1:]:
        xor_arr = xor(prev, block)
        enc_arr = bytearray(aes_f(str(xor_arr)))
        prev = enc_arr
        output.extend(enc_arr)
    return output

def CBC_decrypt(aes_f, blocks):
    """CBC decyrpt: first decrypt then XOR."""
    prev = blocks[0]
    output = bytearray()
    for block in blocks[1:]:
        dec_arr = bytearray(aes_f(str(block)))
        xor_arr = xor(dec_arr, prev)
        prev = block[:]
        output.extend(xor_arr)
    return strip_pkcs7(output)

def apply_CBC(mode, input, key, iv=None, block_len=16):
    """Apply CBC to given input, key, and iv values.
    Args
        mode: str, 'encrypt' or 'decrypt'
        input: bytearray of intput code
        key: str, key to use as AES key
        iv: bytearray, initialization vector for CBC
        block_len: int, optional, normally 16 for block size
    Returns
        Bytearray of encrypted or decrypted input.
    """
    key_AES = AES.new(str(key), AES.MODE_ECB)
    iv_AES = bytearray([0] * block_len) if not iv else iv

    pad = pad_pkcs7 if mode == 'encrypt' else None
    blocks = [iv_AES] + gen_blocks(input, pad, len(key))
    
    return (CBC_encrypt(key_AES.encrypt, blocks) if mode == 'encrypt' else
            CBC_decrypt(key_AES.decrypt, blocks))

# ECB / CBC Encryption and Detection Oracle

def encrypt_ECB_CBC(text, noise=True, force_ECB=False, noise_vals=None,
                    fixed_key=None):
    """Encrypt given text in ECB or CBC, randomly with noise unless specified.
    Args
        text: bytearray of plaintext to encrypt
        noise: bool, add 5 - 10 random before and after text, default True
        force_ECB: bool, force ECB if set to True, default False
        noise_vals: tuple, optional bytearray values to (prepend, append)
        fixed_key: bytearray, optional fixed key to use in encryption
    Returns
        Bytearray of encrypted text.
    """
    if noise and not noise_vals:
        size1 = random.SystemRandom().randint(5, 10)
        size2 = random.SystemRandom().randint(5, 10)
        text = gen_rand_key(size1) + text + gen_rand_key(size2)
    elif noise and noise_vals:
        pre, app = noise_vals
        text = pre + text + app

    # encrypt with ECB or CBC, optionally forcing ECB and fixed key
    key = gen_rand_key() if not fixed_key else fixed_key
    ECB = random.SystemRandom().randint(0, 1)

    if ECB or force_ECB:
        code = apply_ECB('encrypt', text, AES.new(str(key), AES.MODE_ECB))
    else:
        code = apply_CBC('encrypt', text, key, gen_rand_key())
    return code

def gen_ECB_oracle(full_code, rand_prefix=0):