# 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-TBC tweakable block cipher.

This module provides functions to encrypt and decrypt blocks of 128 bits.
"""

from .constants import BLOCK_BYTES, Sbox
from .multiplications import ALPHAS


_permutation = [14, 11, 12, 10, 8, 9, 13, 15, 3, 1, 4, 5, 6, 0, 2, 7]
_permutationInv = [13, 9, 14, 8, 10, 11, 12, 15, 4, 5, 3, 1, 2, 6 ,0 ,7]

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

def _BuildTweakey(tweak, key):
    return tweak+key

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

def _Lane(TK, j):
    return TK[j*8:(j+1)*8]


def _RoundTweakeySchedule(tweakey):
    p = len(tweakey)//8

    multiplied_lanes = (
        ALPHAS[j](_Lane(tweakey, j)) for j in range(p)
    )

    return [byte for lane in multiplied_lanes for byte in lane]


def _SubTweakeyExtract(tweakey, Ci):
    RTKi = [0]*8

    for j, byte in enumerate(tweakey):
        RTKi[j%8] ^= byte

    RTKi[0] ^= Ci

    return RTKi


def _TweakeyScheduleWhole(tweakey, r):
    # Store the initial tweakey in TKs[0], and the corresponding round tweakey
    # in RTKs[0].
    TKs = [tweakey]
    RTKs = [_SubTweakeyExtract(TKs[0], 0)]

    for i in range(1, r):
        TKs.append(_RoundTweakeySchedule(TKs[i-1]))
        RTKs.append(_SubTweakeyExtract(TKs[i], i))

    return RTKs


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

def _NonLinearLayer(state, subtweakey):

    variables_xored = [0 for byte in range(0, 8)]
    for byte in range(0,8):
        variables_xored[byte] = state[byte] ^ subtweakey[byte]

    variables_sboxed = [0 for byte in range(0, 8)]
    for byte in range(0, 8):
        variables_sboxed[byte] = Sbox[variables_xored[byte]]

    state_output = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0,BLOCK_BYTES):
        state_output[byte] = state[byte]
    for byte in range(0, 8):
        state_output[15 - byte] ^= variables_sboxed[byte]

    return state_output


def _LinearLayer(state):
    state_output = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        state_output[byte] = state[byte]

    for byte in range(1, 8):
        state_output[15] ^= state[byte]

    for byte in range(9, 15):
        state_output[byte] ^= state[7]

    return state_output


def _PermutationLayerEnc(state):
    state_output = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        state_output[byte] = state[_permutation[byte]]

    return state_output

def _PermutationLayerDec(state):
    state_output = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        state_output[byte] = state[_permutationInv[byte]]

    return state_output


def _OneRoundEGFNEnc(state, subtweakey):
    state_non_linear = _NonLinearLayer(state, subtweakey)
    state_linear = _LinearLayer(state_non_linear)
    state_permutation = _PermutationLayerEnc(state_linear)

    return state_permutation

def _LastRoundEGFN(state, subtweakey):
    state_non_linear = _NonLinearLayer(state, subtweakey)
    state_linear = _LinearLayer(state_non_linear)

    return state_linear


def _OneRoundEGFNDec(state, subtweakey):
    state_non_linear = _NonLinearLayer(state, subtweakey)
    state_linear = _LinearLayer(state_non_linear)
    state_permutation = _PermutationLayerDec(state_linear)

    return state_permutation


def _Rounds(key_bytes):
    rounds = {
        128: 32,
        192: 36,
        256: 42
    }
    return rounds[key_bytes*8]


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


def encrypt(tweak, key, message):
    r = _Rounds(len(key))

    tweakey = _BuildTweakey(tweak, key)
    RTKs = _TweakeyScheduleWhole(tweakey, r)

    state = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        state[byte] = message[byte]

    for i in range(0, r-1):
        state_output = _OneRoundEGFNEnc(state, RTKs[i])

        for byte in range(0, BLOCK_BYTES):
            state[byte] = state_output[byte]

    state_output = _LastRoundEGFN(state, RTKs[r-1])

    return state_output


def decrypt(tweak, key, cipher):
    r = _Rounds(len(key))

    tweakey = _BuildTweakey(tweak, key)
    RTKs = _TweakeyScheduleWhole(tweakey, r)

    state = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        state[byte] = cipher[byte]

    for i in range(0, r-1):
        state_output = _OneRoundEGFNDec(state, RTKs[r-i-1])

        for byte in range(0, BLOCK_BYTES):
            state[byte] = state_output[byte]

    state_output = _LastRoundEGFN(state, RTKs[0])

    return state_output