//////////////////////////////////////////////////////////////////////////////
/// \file   rijndael.h
/// \brief  Public Domain AES encryption/decryption
/// \author Philip J. Erdelsky <pje@efgh.com>, JD Gascuel, and others.
///
/// \copyright (c) 2015 JD Gascuel. All rights reserved.
/// $Id$
//////////////////////////////////////////////////////////////////////////////
//
// BSD 2-Clause License:
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
//    this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
// THE POSSIBILITY OF SUCH DAMAGE.
//
//////////////////////////////////////////////////////////////////////////////
// HISTORY
//  2002-09-03 PJE: original source http://www.efgh.com/software/rijndael.htm
//  2015-03-14 jDG: Import into OSTC_Companion, major C++ rewrites to insure
//                  key length is not messed up. Added ECb and CFB modes.

#ifndef RIJNDAEL_H
#define RIJNDAEL_H

namespace Rijndael {

#define KEYLENGTH(keybits) ((keybits)/8)
#define RKLENGTH(keybits)  ((keybits)/8+28)
#define NROUNDS(keybits)   ((keybits)/32+6)

typedef unsigned long Word32;
typedef unsigned char Byte;
typedef Byte Block[16];

//////////////////////////////////////////////////////////////////////////////
/// \brief Public Domain AES encryption/decryption rewritten.
///
/// The Rijndael encryption algorithm has been designed to replace the aging
/// DES algorithm. Like DES, it is a block cipher. It uses 128-bit, 192-bit
/// or 256-bit keys. This implementation encrypts 128-bit blocks.
/// (DES used 56-bit keys and 64-bit blocks.)
///
/// The code in this package is a modified version of an implementation
/// placed in the public domain by the following persons:
/// + Vincent Rijmen vincent.rijmen@esat.kuleuven.ac.be
/// + Antoon Bosselaers antoon.bosselaers@esat.kuleuven.ac.be
/// + Paulo Barreto paulo.barreto@terra.com.br
///
/// See details in http://www.efgh.com/software/rijndael.htm
///
struct AES
{
    static void setupEncrypt(Word32* rk, const Byte *key, int keybits);
    static void setupDecrypt(Word32* rk, const Byte *key, int keybits);
    static void encrypt(Word32 *rk, int nrounds, const Block plaintext, Block ciphertext);
    static void decrypt(Word32 *rk, int nrounds, const Block ciphertext, Block plaintext);
    static Word32 get_random(Word32 *rk, int keybits);
};

//////////////////////////////////////////////////////////////////////////////
template<int keybits>
class ECB : private AES
{
    /// \brief Internal state
    ///
    /// storage for encryption buffer, required space:
    /// keybits  32-bit words required
    ///     128      44
    ///     192      52
    ///     256      60
    /// \sa RKLENGTH macro.
    Word32 rk[RKLENGTH(keybits)];

public:
    typedef Byte Key[KEYLENGTH(keybits)];

    ECB();

    //------------------------------------------------------------------------
    /// \brief Initialize encryption state
    ///
    ///
    /// \param[in] key: AES key, where length is:
    ///                 keybits  number of bytes
    ///                     128      16
    ///                     192      24
    ///                     256      32
    ///                 \sa KEYLENGTH macro.
    ///
    /// \code
    ///     Rijndael::ECB<128>::Key key = "my big secret";
    ///     Rijndael::ECB<128> enc;
    ///     enc.setupEncrypt(key);
    ///     ...
    ///     Rijndael::Block plain = "Hello World!";
    ///     Rijndael::Block result;
    ///     enc.encrypt(plain, result);
    /// \endcode
    void setupEncrypt(const Key key);

    //------------------------------------------------------------------------
    /// \brief Encrypt a block of 16 bytes.
    ///
    /// \param[in]   plaintext: The 16 bytes block to encrypt.
    /// \param[out] ciphertext: Space to store the 16 bytes of encrypted data.
    ///
    void encrypt(const Block plaintext, Block ciphertext);

    //------------------------------------------------------------------------
    /// \brief Initialize decryption state.
    ///
    /// \param[in] key: AES key, where length is:
    ///                 keybits  number of bytes
    ///                     128      16
    ///                     192      24
    ///                     256      32
    ///                 \sa KEYLENGTH macro.
    ///
    /// \code
    ///     Rijndael::ECB<128>::Key key = "my big secret";
    ///     Rijndael::ECB<128> dec;
    ///     dec.setupDecrypt(key);
    ///     ...
    ///     Rijndael::Block cipher = ...;
    ///     Rijndael::Block result;
    ///     Rijndael::decrypt(cipher, result);
    /// \endcode
    void setupDecrypt(const Key key);

    //------------------------------------------------------------------------
    /// \brief Decrypt a block of 16 bytes.
    ///
    /// \param[in]  ciphertext: The 16 bytes block of data to decrypt.
    /// \param[out]  plaintext: Space to store the 16 bytes result block.
    ///
    void decrypt(const Block ciphertext, Block plaintext);

    //----------------------------------------------------------------------------
    /// \brief Crypto base PRNG
    ///
    /// Based on wall-clock value and current key, but should be a crypto-secure
    /// generator.
    Word32 get_random();
};

//////////////////////////////////////////////////////////////////////////////
template<int keybits>
class CFB : public ECB<keybits>
{
    /// Initialization vector (salt), updated in CFB mode
    /// for the next block of text.
    Block iv;

public:
    typedef Byte Key[KEYLENGTH(keybits)];
    typedef Block IV;

    //------------------------------------------------------------------------
    /// \brief Initialize encryption/decription state
    ///
    ///
    /// \param[in] key: AES key, where length is:
    ///                 keybits  number of bytes
    ///                     128      16
    ///                     192      24
    ///                     256      32
    ///                 \sa KEYLENGTH macro.
    /// \param[in]  iv: initialization vector. Some randomness needed to
    ///                 enforce the sequence is non replayable.
    ///
    /// \code
    ///     Rijndael::CFB<128>::Key key = "my big secret";
    ///     Rijndael::CFB<128> enc(key, iv);
    ///     ...
    ///     Rijndael::Block plain = "Hello World!";
    ///     Rijndael::Block cipher, again;
    ///     enc.encrypt(plain, cipher);
    ///     enc.decrypt(cipher, again);
    /// \endcode
    CFB(const Key key, const IV iv);

    //------------------------------------------------------------------------
    /// \brief Encrypt a block of 16 bytes, with IV (CFB mode)
    ///
    /// \param[in]   plaintext: The 16 bytes block to encrypt.
    /// \param[out] ciphertext: Space to store the 16 bytes of encrypted data.
    ///
    void encrypt(const Block plaintext, Block ciphertext);

    //----------------------------------------------------------------------------
    /// \brief Decrypt a block of 16 bytes, with IV (CFB mode)
    ///
    /// \param[in]  ciphertext: The 16 bytes block of data to decrypt.
    /// \param[out]  plaintext: Space to store the 16 bytes result block.
    ///
    void decrypt(const Block ciphertext, Block plaintext);
};

//////////////////////////////////////////////////////////////////////////////

template<int keybits>
ECB<keybits>::ECB()
{
    for(int i=0; i<RKLENGTH(keybits); ++i)
        rk[i] = 0;
}

template<int keybits>
void ECB<keybits>::setupEncrypt(const Key key)
{
    AES::setupEncrypt(rk, key, keybits);
}

template<int keybits>
void ECB<keybits>::encrypt(const Block plaintext, Block ciphertext)
{
    AES::encrypt(rk, NROUNDS(keybits), plaintext, ciphertext);
}

//////////////////////////////////////////////////////////////////////////////

template<int keybits>
void ECB<keybits>::setupDecrypt(const Key key)
{
    AES::setupDecrypt(rk, key, keybits);
}

template<int keybits>
void ECB<keybits>::decrypt(const Block ciphertext, Block plaintext)
{
    AES::decrypt(rk, NROUNDS(keybits), ciphertext, plaintext);
}

//////////////////////////////////////////////////////////////////////////////

template<int keybits>
Word32 ECB<keybits>::get_random()
{
    return AES::get_random(rk, keybits);
}

//////////////////////////////////////////////////////////////////////////////

template<int keybits>
CFB<keybits>::CFB(const Key key, const IV _iv)
{
    ECB<keybits>::setupEncrypt(key);
    for(int i=0; i<16; ++i)
        iv[i] = _iv[i];
}

template<int keybits>
void CFB<keybits>::encrypt(const Block plaintext, Block ciphertext)
{
    Block tmp;
    ECB<keybits>::encrypt(iv, tmp);
    for(int i=0; i<16; ++i)
    {
        ciphertext[i] = plaintext[i] ^ tmp[i];
        iv[i]         = ciphertext[i];
    }
}

template<int keybits>
void CFB<keybits>::decrypt(const Block ciphertext, Block plaintext)
{
    Block tmp;
    ECB<keybits>::encrypt(iv, tmp);
    for(int i=0; i<16; ++i)
    {
        plaintext[i] = ciphertext[i] ^ tmp[i];
        iv[i]        = ciphertext[i];
    }
}

}
#endif // RIJNDAEL_H