////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  CryptLib_Aes
//
//  Implementation of AES block cipher. Originally written by Kokke (https://github.com/kokke). Modified by WaterJuice
//  retaining Public Domain license.
//
//  AES is a block cipher that operates on 128 bit blocks. Encryption an Decryption routines use an AesContext which
//  must be initialised with the key. An AesContext can be initialised with a 128, 192, or 256 bit key. Use the
//  AesInitialise[n] functions to initialise the context with the key. Once an AES context is initialised its contents
//  are not changed by the encrypting and decrypting functions. A context only needs to be initialised once for any
//  given key and the context may be used by the encrypt/decrypt functions in simultaneous threads.
//  All operations are performed byte wise and this implementation works in both little and endian processors.
//  There are no alignment requirements with the keys and data blocks.
//
//  This is free and unencumbered software released into the public domain - November 2017 waterjuice.org
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#pragma once

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  IMPORTS
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <stdint.h>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  TYPES
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#define AES_KEY_SIZE_128        16
#define AES_KEY_SIZE_192        24
#define AES_KEY_SIZE_256        32
#define AES_BLOCK_SIZE          16

// AesContext - This must be initialised using AesInitialise128, AesInitialise192 or AesInitialise256
// Do not modify the contents of this structure directly.
typedef struct
{
    uint32_t    KeySizeInWords;
    uint32_t    NumberOfRounds;
    uint8_t     RoundKey[240];
} AesContext;

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  PUBLIC FUNCTIONS
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesInitialise128
//
//  Initialises an AesContext with a 128 bit key.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesInitialise128
    (
        uint8_t const   Key [AES_KEY_SIZE_128],         // [in]
        AesContext*     Context                         // [out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesInitialise192
//
//  Initialises an AesContext with a 192 bit key.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesInitialise192
    (
        uint8_t const   Key [AES_KEY_SIZE_192],         // [in]
        AesContext*     Context                         // [out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesInitialise256
//
//  Initialises an AesContext with a 256 bit key.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesInitialise256
    (
        uint8_t const   Key [AES_KEY_SIZE_256],         // [in]
        AesContext*     Context                         // [out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesEncrypt
//
//  Performs an AES encryption of one block (128 bits) with the AesContext initialised with one of the functions
//  AesInitialise[n]. Input and Output can point to same memory location, however it is more efficient to use
//  AesEncryptInPlace in this situation.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesEncrypt
    (
        AesContext const*   Context,                    // [in]
        uint8_t const       Input [AES_BLOCK_SIZE],     // [in]
        uint8_t             Output [AES_BLOCK_SIZE]     // [out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesDecrypt
//
//  Performs an AES decryption of one block (128 bits) with the AesContext initialised with one of the functions
//  AesInitialise[n]. Input and Output can point to same memory location, however it is more efficient to use
//  AesDecryptInPlace in this situation.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesDecrypt
    (
        AesContext const*   Context,                    // [in]
        uint8_t const       Input [AES_BLOCK_SIZE],     // [in]
        uint8_t             Output [AES_BLOCK_SIZE]     // [out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesEncryptInPlace
//
//  Performs an AES encryption of one block (128 bits) with the AesContext initialised with one of the functions
//  AesInitialise[n]. The encryption is performed in place.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesEncryptInPlace
    (
        AesContext const*   Context,                    // [in]
        uint8_t             Block [AES_BLOCK_SIZE]      // [in out]
    );

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//  AesDecryptInPlace
//
//  Performs an AES decryption of one block (128 bits) with the AesContext initialised with one of the functions
//  AesInitialise[n]. The decryption is performed in place.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
    AesDecryptInPlace
    (
        AesContext const*   Context,                    // [in]
        uint8_t             Block [AES_BLOCK_SIZE]      // [in out]
    );