# Implementation of the Lilliput-AE tweakable block cipher.
#
# Authors, hereby denoted as "the implementer":
#     Kévin Le Gouguec,
#     Léo Reynaud
#     2019.
#
# For more information, feedback or questions, refer to our website:
# https://paclido.fr/lilliput-ae
#
# To the extent possible under law, the implementer has waived all copyright
# and related or neighboring rights to the source code in this file.
# http://creativecommons.org/publicdomain/zero/1.0/

"""Lilliput-I Authenticated Encryption mode.

This module provides the functions for authenticated encryption and decryption
using Lilliput-AE's nonce-respecting mode based on ΘCB3.
"""

from enum import Enum

from .constants import BLOCK_BYTES, NONCE_BYTES
from .ae_common import (
    bytes_to_block_matrix,
    block_matrix_to_bytes,
    build_auth,
    pad10,
    TagValidationError,
    xor
)
from . import tbc


TWEAK_BITS = 192
TWEAK_BYTES = TWEAK_BITS//8


def _low_part(array, number_bits):
    shifted = 0
    for byte in range(0, len(array)):
        shifted |= (array[byte] << (8 * byte))

    mask = 0
    for bit in range(0, number_bits):
        mask |= (0x1 << bit)

    lower_part = shifted & mask

    will_pad = 0
    if number_bits % 8 != 0:
        will_pad = 1

    lower_part_byte = []
    nb_bytes = number_bits//8 + will_pad
    for byte in range(nb_bytes):
        lower_part_byte.append(lower_part & 0xff)
        lower_part = lower_part >> 8

    return lower_part_byte


class _MessageTweak(Enum):
    BLOCK = 0b0000
    NO_PADDING = 0b0001
    PAD = 0b0100
    FINAL = 0b0101


def _tweak_message(N, j, padding):
    tweak = [0 for byte in range(0, TWEAK_BYTES)]
    for byte in range(NONCE_BYTES-1, -1, -1):
        tweak[byte + (TWEAK_BYTES-NONCE_BYTES)] |= (N[byte] & 0xf0) >> 4
        tweak[byte + (TWEAK_BYTES-NONCE_BYTES-1)] |= (N[byte] & 0x0f) << 4

    tweak[TWEAK_BYTES-NONCE_BYTES-1] |= ((j >> 64) & 0xf)
    for byte in range(TWEAK_BYTES-NONCE_BYTES-2, -1, -1):
        tweak[byte] = (j >> (8 * byte)) & 0xff

    tweak[-1] |= padding.value<<4

    return tweak


def _treat_message_enc(M, N, key):
    checksum = [0 for byte in range(0, BLOCK_BYTES)]

    l = len(M)//BLOCK_BYTES
    padding_bytes = len(M)%BLOCK_BYTES

    M = bytes_to_block_matrix(M)
    C = []

    for j in range(0, l):
        checksum = xor(checksum, M[j])
        tweak = _tweak_message(N, j, _MessageTweak.BLOCK)
        C.append(tbc.encrypt(tweak, key, M[j]))

    if padding_bytes == 0:
        tweak = _tweak_message(N, l, _MessageTweak.NO_PADDING)
        Final = tbc.encrypt(tweak, key, checksum)

    else:
        m_padded = pad10(M[l])
        checksum = xor(checksum, m_padded)
        tweak = _tweak_message(N, l, _MessageTweak.PAD)
        pad = tbc.encrypt(tweak, key, [0 for byte in range(0, BLOCK_BYTES)])

        lower_part = _low_part(pad, padding_bytes*8)
        C.append(xor(M[l], lower_part))
        tweak_final = _tweak_message(N, l+1, _MessageTweak.FINAL)
        Final = tbc.encrypt(tweak_final, key, checksum)

    return (Final, C)


def _treat_message_dec(C, N, key):
    checksum = [0 for byte in range(0, BLOCK_BYTES)]

    l = len(C)//BLOCK_BYTES
    padding_bytes = len(C)%BLOCK_BYTES

    C = bytes_to_block_matrix(C)
    M = []

    for j in range(0, l):
        tweak = _tweak_message(N, j, _MessageTweak.BLOCK)
        M.append(tbc.decrypt(tweak, key, C[j]))
        checksum = xor(checksum, M[j])

    if padding_bytes == 0:
        tweak = _tweak_message(N, l, _MessageTweak.NO_PADDING)
        Final = tbc.encrypt(tweak, key, checksum)

    else:
        tweak = _tweak_message(N, l, _MessageTweak.PAD)
        pad = tbc.encrypt(tweak, key, [0 for byte in range(0, BLOCK_BYTES)])
        lower_part = _low_part(pad, padding_bytes*8)
        M.append(xor(C[l], lower_part))

        m_padded = pad10(M[l])
        checksum = xor(checksum, m_padded)
        tweak_final = _tweak_message(N, l+1, _MessageTweak.FINAL)
        Final = tbc.encrypt(tweak_final, key, checksum)

    return (Final, M)


def encrypt(A, M, N, key):
    K = list(key)

    Auth = build_auth(TWEAK_BITS, A, K)
    (Final, C) = _treat_message_enc(M, N, K)
    tag = xor(Auth, Final)

    return block_matrix_to_bytes(C), bytes(tag)


def decrypt(A, C, N, tag, key):
    K = list(key)
    tag = list(tag)

    Auth = build_auth(TWEAK_BITS, A, K)
    (Final, M) = _treat_message_dec(C, N, K)
    tag2 = xor(Auth, Final)

    if tag != tag2:
        raise TagValidationError(tag, tag2)

    return block_matrix_to_bytes(M)