import lilliput_ae_1
import lilliput_ae_2

BLOCK_BYTES = 16
N_BYTES = 15

def _getParameters(mode=1, key_length=128) :
    rounds = {
        128: 32,
        192: 36,
        256: 42
    }

    tweak_lengths = {
        1: 192,
        2: 128
    }

    return tweak_lengths[mode], rounds[key_length]


def ArrayToBlockbytesMatrix(array) :
    length = len(array)
    pad = 0
    if(length % BLOCK_BYTES == 0) :
        number_blocks = int(length / BLOCK_BYTES)
    else :
        number_blocks = int((length + (BLOCK_BYTES - (length % BLOCK_BYTES))) / BLOCK_BYTES)
        pad = 1

    matrix = [[0] * BLOCK_BYTES for block in range(0, number_blocks - pad)]
    if(pad == 1) :
        matrix.append([0] * (length % BLOCK_BYTES))

    for byte in range(0, length) :
        matrix[int(byte / BLOCK_BYTES)][byte % BLOCK_BYTES] = array[byte]

    return matrix

def BlockbytesMatrixToBytes(matrix):
    return bytes(byte for block in matrix for byte in block)

############################################

def _checkInputs(key, nonce):
    valid_key_lengths = (128, 192, 256)

    if len(key)*8 not in valid_key_lengths:
        raise ValueError('invalid key size: {} not in {}'.format(len(key)*8, valid_key_lengths))

    if len(nonce) != N_BYTES:
        raise ValueError('nonce must be {}-byte long'.format(N_BYTES))


def mainEnc(plaintext, adata, key, nonce, mode):
    _checkInputs(key, nonce)

    key_bits = len(key)*8
    tweak_bits, rounds = _getParameters(mode, key_bits)

    A = adata
    M = plaintext
    N = nonce

    A_BITS = 8 * len(A)
    M_BITS = 8 * len(M)

    A = ArrayToBlockbytesMatrix(A)
    M = ArrayToBlockbytesMatrix(M)

    if(mode == 1) :
        (C, tag) = lilliput_ae_1.OCB3Enc(A, M, N, A_BITS, M_BITS, key, key_bits, tweak_bits, rounds)
    if(mode == 2) :
        (C, tag) = lilliput_ae_2.SCT2Enc(A, M, N, A_BITS, M_BITS, key, key_bits, tweak_bits, rounds)

    return BlockbytesMatrixToBytes(C), bytes(tag)


def mainDec(ciphertext, tag, adata, key, nonce, mode):
    _checkInputs(key, nonce)

    key_bits = len(key)*8
    tweak_bits, rounds = _getParameters(mode, key_bits)

    A = adata
    C = ciphertext
    N = nonce
    tag = list(tag)

    M_BITS = 8 * len(C)
    A_BITS = 8 * len(A)

    A = ArrayToBlockbytesMatrix(A)
    C = ArrayToBlockbytesMatrix(C)

    if(mode == 1) :
        M = lilliput_ae_1.OCB3Dec(A, C, N, tag, A_BITS, M_BITS, key, key_bits, tweak_bits, rounds)
    if(mode == 2) :
        M = lilliput_ae_2.SCT2Dec(A, C, N, tag, A_BITS, M_BITS, key, key_bits, tweak_bits, rounds)

    return BlockbytesMatrixToBytes(M)