/*
 * Copyright (C) 2020-2021 Aloysius Indrayanto
 *
 * This file is part of the eCxx library, see LICENSE file for the license details.
 */

/*
 * ----------------------------------------------------------------------------------
 * The code in this file is copied (with a slight modifications) from:
 *     Crypto.cpp
 * ----------------------------------------------------------------------------------
 * BearSSL Copyright (C) 2016 Thomas Pornin (pornin@bolet.org)
 * The rest of this file Copyright (C) 2019 Anders Löfgren
 *
 * MIT License
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in the
 * Software without restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
 * Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
 * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 * ----------------------------------------------------------------------------------
 */


#include <assert.h>
#include <bearssl/bearssl.h>

#include "Crypto.h"


namespace eCxx {

namespace NodeMCU {

namespace Crypto {


//
// Helper functions
//
namespace TypeCast = NodeMCU::TypeConversion;


static uint8_t* defaultNonceGenerator(uint8_t* nonceArray, const size_t nonceLength)
{ return random(nonceArray, nonceLength); }


size_t             _ctMinDataLength =    0;
size_t             _ctMaxDataLength = 1024;
NonceGeneratorType _nonceGenerator = defaultNonceGenerator;


static void* createBearsslHmac(const br_hash_class* hashType, const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{
    br_hmac_key_context keyContext;
    br_hmac_context     hmacContext;

    br_hmac_key_init(&keyContext, hashType, hashKey, hashKeyLength);
    br_hmac_init(&hmacContext, &keyContext, outputLength);
    br_hmac_update(&hmacContext, data, dataLength);
    br_hmac_out(&hmacContext, resultArray);

    return resultArray;
}


static const WString createBearsslHmac(const br_hash_class* hashType, const uint8_t hashTypeNaturalLength, const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{
    eCxx::ignoreUnused(hashTypeNaturalLength);

    assert(1 <= hmacLength && hmacLength <= hashTypeNaturalLength);

    uint8_t hmac[hmacLength];
    createBearsslHmac(hashType, message.c_str(), message.length(), hashKey, hashKeyLength, hmac, hmacLength);

    return TypeCast::uint8ArrayToHexString(hmac, hmacLength);
}


static void* createBearsslHmacCT(const br_hash_class* hashType, const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{
    assert( _ctMinDataLength <= dataLength && dataLength <= _ctMaxDataLength );

    br_hmac_key_context keyContext;
    br_hmac_context     hmacContext;

    br_hmac_key_init(&keyContext, hashType, hashKey, hashKeyLength);
    br_hmac_init(&hmacContext, &keyContext, outputLength);
    br_hmac_outCT(&hmacContext, data, dataLength, _ctMinDataLength, _ctMaxDataLength, resultArray);

    return resultArray;
}


static const WString createBearsslHmacCT(const br_hash_class* hashType, const uint8_t hashTypeNaturalLength, const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{
    eCxx::ignoreUnused(hashTypeNaturalLength);

    assert( 1 <= hmacLength && hmacLength <= hashTypeNaturalLength );

    uint8_t hmac[hmacLength];
    createBearsslHmacCT(hashType, message.c_str(), message.length(), hashKey, hashKeyLength, hmac, hmacLength);

    return TypeCast::uint8ArrayToHexString(hmac, hmacLength);
}


static void* md5HashHelper(const void* data, const size_t dataLength, void* resultArray)
{
    br_md5_context context;
    br_md5_init(&context);
    br_md5_update(&context, data, dataLength);
    br_md5_out(&context, resultArray);

    return resultArray;
}


static void* sha1HashHelper(const void* data, const size_t dataLength, void* resultArray)
{
    br_sha1_context context;
    br_sha1_init(&context);
    br_sha1_update(&context, data, dataLength);
    br_sha1_out(&context, resultArray);

    return resultArray;
}


//
// Exported functions
//
extern void setCtMinDataLength(const size_t ctMinDataLength)
{
    assert( getCtMaxDataLength() - ctMinDataLength <= CT_MAX_DIFF );
    _ctMinDataLength = ctMinDataLength;
}


extern size_t getCtMinDataLength()
{ return _ctMinDataLength; }


extern void setCtMaxDataLength(const size_t ctMaxDataLength)
{
    assert( ctMaxDataLength - getCtMinDataLength() <= CT_MAX_DIFF );
    _ctMaxDataLength = ctMaxDataLength;
}


extern size_t getCtMaxDataLength()
{ return _ctMaxDataLength; }


extern void setNonceGenerator(NonceGeneratorType nonceGenerator)
{ _nonceGenerator = nonceGenerator; }


extern NonceGeneratorType getNonceGenerator()
{ return _nonceGenerator; }


//
// Implementation of classes
//
void* MD5::hash(const void* data, const size_t dataLength, void* resultArray)
{ return md5HashHelper(data, dataLength, resultArray); }


const WString MD5::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    md5HashHelper(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* MD5::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_md5_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString MD5::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_md5_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* MD5::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_md5_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString MD5::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_md5_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* SHA1::hash(const void* data, const size_t dataLength, void* resultArray)
{ return sha1HashHelper(data, dataLength, resultArray); }


const WString SHA1::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    sha1HashHelper(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* SHA1::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_sha1_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA1::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_sha1_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* SHA1::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_sha1_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA1::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_sha1_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* SHA224::hash(const void* data, const size_t dataLength, void* resultArray)
{
    br_sha224_context context;
    br_sha224_init(&context);
    br_sha224_update(&context, data, dataLength);
    br_sha224_out(&context, resultArray);

    return resultArray;
}


const WString SHA224::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    hash(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* SHA224::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_sha224_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA224::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_sha224_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* SHA224::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_sha224_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA224::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_sha224_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* SHA256::hash(const void* data, const size_t dataLength, void* resultArray)
{
    br_sha256_context context;
    br_sha256_init(&context);
    br_sha256_update(&context, data, dataLength);
    br_sha256_out(&context, resultArray);

    return resultArray;
}


const WString SHA256::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    hash(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* SHA256::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_sha256_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }

const WString SHA256::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_sha256_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* SHA256::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_sha256_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA256::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_sha256_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* SHA384::hash(const void* data, const size_t dataLength, void* resultArray)
{
    br_sha384_context context;
    br_sha384_init(&context);
    br_sha384_update(&context, data, dataLength);
    br_sha384_out(&context, resultArray);

    return resultArray;
}


const WString SHA384::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    hash(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* SHA384::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_sha384_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA384::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_sha384_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* SHA384::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_sha384_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA384::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_sha384_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* SHA512::hash(const void* data, const size_t dataLength, void* resultArray)
{
    br_sha512_context context;
    br_sha512_init(&context);
    br_sha512_update(&context, data, dataLength);
    br_sha512_out(&context, resultArray);

    return resultArray;
}


const WString SHA512::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    hash(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


void* SHA512::hmac(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmac(&br_sha512_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA512::hmac(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmac(&br_sha512_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


void* SHA512::hmacCT(const void* data, const size_t dataLength, const void* hashKey, const size_t hashKeyLength, void* resultArray, const size_t outputLength)
{ return createBearsslHmacCT(&br_sha512_vtable, data, dataLength, hashKey, hashKeyLength, resultArray, outputLength); }


const WString SHA512::hmacCT(const WString& message, const void* hashKey, const size_t hashKeyLength, const size_t hmacLength)
{ return createBearsslHmacCT(&br_sha512_vtable, NATURAL_LENGTH, message, hashKey, hashKeyLength, hmacLength); }


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


void* MD5SHA1::hash(const void* data, const size_t dataLength, void* resultArray)
{
    br_md5sha1_context context;
    br_md5sha1_init(&context);
    br_md5sha1_update(&context, data, dataLength);
    br_md5sha1_out(&context, resultArray);

    return resultArray;
}


const WString MD5SHA1::hash(const WString& message)
{
    uint8_t hashArray[NATURAL_LENGTH];
    hash(message.c_str(), message.length(), hashArray);

    return TypeCast::uint8ArrayToHexString(hashArray, NATURAL_LENGTH);
}


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


HKDF::HKDF(const void* keyMaterial, const size_t keyMaterialLength, const void* salt, const size_t saltLength)
{ init(keyMaterial, keyMaterialLength, salt, saltLength); }


void HKDF::init(const void* keyMaterial, const size_t keyMaterialLength, const void* salt, const size_t saltLength)
{
    br_hkdf_init(&hkdfContext, &br_sha256_vtable, salt, saltLength);
    br_hkdf_inject(&hkdfContext, keyMaterial, keyMaterialLength);
    br_hkdf_flip(&hkdfContext);
}


size_t HKDF::produce(void* resultArray, const size_t outputLength, const void* info, const size_t infoLength)
{ return br_hkdf_produce(&hkdfContext, info, infoLength, resultArray, outputLength); }


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


void chacha20Poly1305Kernel(const int encrypt, void* data, const size_t dataLength, const void* key, const void* keySalt, const size_t keySaltLength, const void* nonce, void* tag, const void* aad, const size_t aadLength)
{
    if(keySalt == nullptr) {
        br_poly1305_ctmul32_run(key, nonce, data, dataLength, aad, aadLength, tag, br_chacha20_ct_run, encrypt);
    }
    else {
        HKDF    hkdfInstance(key, ENCRYPTION_KEY_LENGTH, keySalt, keySaltLength);
        uint8_t derivedEncryptionKey[ENCRYPTION_KEY_LENGTH] {0};

        hkdfInstance.produce(derivedEncryptionKey, ENCRYPTION_KEY_LENGTH);
        br_poly1305_ctmul32_run(derivedEncryptionKey, nonce, data, dataLength, aad, aadLength, tag, br_chacha20_ct_run, encrypt);
    }
}


void ChaCha20Poly1305::encrypt(void* data, const size_t dataLength, const void* key, const void* keySalt, const size_t keySaltLength, void* resultingNonce, void* resultingTag, const void* aad, const size_t aadLength)
{
    uint8_t* nonce = reinterpret_cast<uint8_t*>(resultingNonce);
    getNonceGenerator()(nonce, 12);

    chacha20Poly1305Kernel(1, data, dataLength, key, keySalt, keySaltLength, nonce, resultingTag, aad, aadLength);
}


bool ChaCha20Poly1305::decrypt(void* data, const size_t dataLength, const void* key, const void* keySalt, const size_t keySaltLength, const void* encryptionNonce, const void* encryptionTag, const void* aad, const size_t aadLength)
{
    const uint8_t* oldTag = reinterpret_cast<const uint8_t*>(encryptionTag);
          uint8_t  newTag[16] {0};

    chacha20Poly1305Kernel(0, data, dataLength, key, keySalt, keySaltLength, encryptionNonce, newTag, aad, aadLength);

    for(uint32_t i = 0; i < sizeof newTag; ++i) {
        if(newTag[i] != oldTag[i]) return false;
    }

    return true;
}


} // namespace Crypto

} // namespace NodeMCU

} // namespace eCxx
