/*!
 * \brief   This provides device independent RSA 2048 bit data signing.
 *
 * See the header file for details about this module.
 *
 * \copyright   2023, Sunrise Labs, Inc.
 *
 * \file   sign_rsa.cpp
 *
 * \author PBraica
 * \date   March 2023
 */

#include "SignRsa.h"
#include <cstring>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>

// Windows requires this, Linux doesn't have it.
#if defined(_WIN32) || defined(WIN32)
#include <openssl/applink.c>
#endif

/**
 * This class pulls some of the RSA logic into one place that 
 * could be error prone, and it wraps the kind and size of the hash
 * currently as SHA256 so that in the future other hashes or 
 * custom hashes could be used.
 */
class _InternalDigest {
public:
    /*!
     * \brief   Constructor.
     */
    _InternalDigest():
        _ctx(),
        _pDigest(),
        _finalized(false) 
    {
        SHA256_Init(&_ctx);
    }

    /*!
     * \brief   Copy from other.
     * 
     * \param   pOther      Other object.
     */
    void copy_from(const _InternalDigest* pOther) {
        // CTX is a structure in RSA.
        memcpy(&_ctx, &pOther->_ctx, sizeof(SHA256_CTX));

        // Copy the digest.
        memcpy(_pDigest, pOther->_pDigest, SHA256_DIGEST_LENGTH);
    }

    /*!
     * \brief   Update the context. 
     * 
     * \param   pData       Buffer pointer.
     * \param   size_bytes   Size in bytes.
     */
    void update(const unsigned char * pData, std::size_t size_bytes) {
        if (size_bytes > 0) {
            SHA256_Update(&_ctx, pData, size_bytes);
        }
    }

    /*!
     * \brief   Finalize the context and create the digest bytes.
     * 
     * \return  Digest bytes.
     */
    unsigned char* get_digest() {
        SHA256_Final(_pDigest, &_ctx);
        return _pDigest;
    }

    /*!
     * \brief   Clear.
     */
    void clear() {
        SHA256_Init(&_ctx);
    }

protected:
    /*!
     * \brief   The context computation object.
     */
    SHA256_CTX      _ctx;   

    /*!
     * \brief   The Digest bytes.
     */
    unsigned char   _pDigest[SHA256_DIGEST_LENGTH];

    /*! \brief  Was this finalized? */
    bool            _finalized;
};

// Buffer for file read operations. The buffer must be able to accomodate
// the RSA signature in whole (e.g. 4096-bit RSA key produces 512 byte signature)
#define BUFFER_SIZE (16384)

/*!
 * \brief    Do the initialization calls for openSSL.
 */
class DoOnce {
public:
    /*!
     * \brief   Constructor.
     */
    DoOnce() {            
        OpenSSL_add_all_algorithms();
        ERR_load_crypto_strings();
    }
};

/*!
 * \brief   Constructor.
 */
SignRsa::SignRsa():
    _pDigest(new _InternalDigest()),
    _pPublic(nullptr),
    _pPrivate(nullptr) {
    
    // On first entry do this.
    static DoOnce doOnce;
}

/*!
 * \brief   Copy constructor.
 */
SignRsa::SignRsa(const SignRsa& other) :
    _pDigest(new _InternalDigest()),
    _pPublic(nullptr),
    _pPrivate(nullptr),
    _lastPub("")
{
    _pDigest->copy_from(other._pDigest);

    _pPublic = (other._pPublic == nullptr) ? 
        nullptr : RSAPublicKey_dup(other._pPublic);
    _pPrivate = (other._pPrivate == nullptr) ? 
        nullptr : RSAPrivateKey_dup(other._pPrivate);
}

/*!
 * \brief   Destructor.
 */
SignRsa::~SignRsa() {
    if (_pDigest != nullptr) {
        delete _pDigest;
    }
    if (_pPublic != nullptr) {
        RSA_free(_pPublic);
    }
    if (_pPrivate != nullptr) {
        RSA_free(_pPrivate);
    }
}

/*!
 * \brief   Create new RSA keys.
 */
void SignRsa::create_keys() {

    // Free old if needed.
    if (_pPublic != nullptr) {
        RSA_free(_pPublic);
    }
    if (_pPrivate != nullptr) {
        RSA_free(_pPrivate);
    }

    // Create new.
    unsigned long e = RSA_F4;
    BIGNUM *bne = BN_new();
    int ret = BN_set_word(bne, e);
    if (ret == 1) {
        _pPrivate = RSA_new();
        ret = RSA_generate_key_ex(_pPrivate, 2048, bne, NULL);
        if (ret == 1) {
            _pPublic = RSAPublicKey_dup(_pPrivate);
        } else {
            RSA_free(_pPrivate);
        }
    } 
    BN_free(bne);
    return;    
}

/*!
 * \brief   Save this objects keys to the folder if they exist.
 *
 * Saves them as public.pem and private.pem, if they exist.
 * 
 * \param   folder_name     Folder to use.
 */
void SignRsa::save_keys(const std::string& folder_name) {

    const std::string pubName = folder_name + "/public.pem";
    const std::string priName = folder_name + "/private.pem";

    if (_pPublic != nullptr) {
        FILE* pemFile = fopen(pubName.c_str(), "wb");
        if (pemFile != NULL) {
            PEM_write_RSA_PUBKEY(pemFile, _pPublic);
            fclose(pemFile);
        }
    }
    if (_pPrivate != nullptr) {
        FILE* pemFile = fopen(priName.c_str(), "wb");
        if (pemFile != NULL) {
            PEM_write_RSAPrivateKey(pemFile, _pPrivate, NULL, NULL, 0, NULL, NULL);
            fclose(pemFile);
        }
    }
}

/*!
 * \brief   Load the public.pem and private.pem, what ever exists.
 *
 * \param   folder_name     Folder to use, or the path+name of the public or private key.
 */
void SignRsa::load_keys(const std::string& folder_name) {

    std::string key_loc = folder_name;
    std::size_t clip = key_loc.find("public.pem");
    if (clip != std::string::npos) {
        key_loc.resize(clip - 1);
    }
    clip = key_loc.find("private.pem");
    if (clip != std::string::npos) {
        key_loc.resize(clip - 1);
    }

    const std::string pubName = key_loc + "/public.pem";
    const std::string priName = key_loc + "/private.pem";

    if (_pPublic != nullptr) {
        RSA_free(_pPublic);
        _pPublic = nullptr;
    }
    if (_pPrivate != nullptr) {
        RSA_free(_pPrivate);
        _pPrivate = nullptr;
    }

    FILE* pPubFile = NULL;
    FILE* pPriFile = NULL;
    std::string tmpString="";

    // Try the _pPublic.
    try {
        pPubFile = fopen(pubName.c_str(), "r");
        if (pPubFile != NULL) {
            // Get size.
            fseek(pPubFile, 0, SEEK_END);
            int size = ftell(pPubFile);

            // Got back and read into tmpString.
            fseek(pPubFile, 0, SEEK_SET);
            tmpString.resize(size, '\0');

            // Read to cache the string value.
            std::size_t readBytes = fread(&tmpString[0], sizeof(char), (size_t)size, pPubFile);
            if (readBytes != (std::size_t)size) {
                tmpString = "";
            }
            
            // Go back and read public key from file
            fseek(pPubFile, 0, SEEK_SET);
            _pPublic = PEM_read_RSA_PUBKEY(pPubFile, NULL, NULL, NULL);
        }
    } catch (...) {
        _pPublic = nullptr;
    }

    // Cache the public key since it was read ok, 
    // this helps when validating many files from the same API.
    if (_pPublic != nullptr) {
        _lastPub = tmpString;
    }

    // Try the _pPrivate.
    try {
        pPriFile = fopen(priName.c_str(), "r");
        if (pPriFile != NULL) {
            // Read public key from file
            _pPrivate = PEM_read_RSAPrivateKey(pPriFile, NULL, NULL, NULL);
        }
    }
    catch (...) {
        _pPrivate = nullptr;
    }

    // Because these are typically security procedures,
    // if they fail we don't want to crash, we need to be able
    // to keep going so the close is wrapped in try/catch.

    try 
    {
        if (pPubFile) 
        {
            fclose(pPubFile);
        }
    }
    catch (...) {
        ; // NOP.
    }

    try
    {
        if (pPriFile)
        {
            fclose(pPriFile);
        }
    }
    catch (...) {
        ; // NOP.
    }
}

/*!
 * \brief   Set the public key from a string.
 *
 * \note    We cache public keys for speed / ease of use.
 * 
 * \param   key      The key, base 64.
 * 
 * \return  True on success.
 */
bool SignRsa::set_public(const std::string& key) {

    if (key == _lastPub) 
    {
        return true; // No need to do anything and thus success.
    }

    if (_pPublic != nullptr) {
        RSA_free(_pPublic);
    }

    BIO* keybio;
    const char* cstr = key.c_str();
    keybio = BIO_new_mem_buf((void*)cstr, -1);
    if (keybio == NULL) 
    {
        return false;
    }
    _pPublic = PEM_read_bio_RSA_PUBKEY(keybio, &_pPublic, NULL, NULL);
    BIO_free(keybio);

    _lastPub = key;
    return true;
}

/*!
 * \brief   Set the private key from a string.
 *
 * \note    Private keys are NOT cached as text for security!
 * 
 * \param   key      The key, base 64
 *
 * \return  True on success.
 */
bool SignRsa::set_private(const std::string& key) {
    if (_pPrivate != nullptr) {
        RSA_free(_pPrivate);
    }
    BIO* keybio;
    const char* cstr = key.c_str();
    keybio = BIO_new_mem_buf((void*)cstr, -1);
    if (keybio == NULL) {
        return false;
    }
    _pPrivate = PEM_read_bio_RSAPrivateKey(keybio, &_pPrivate, NULL, NULL);
    return true;
}

/*!
 * \brief   Clear our digest to do a different new stream.
 */
void SignRsa::clear_digest() {
    _pDigest->clear();
}

/*!
 * \brief   Update our Digest object from a buffer.
 *
 * \param   pData       Data pointer.
 * \param   size_bytes  Size in bytes.
 */
void SignRsa::update_digest(const unsigned char* pData, std::size_t size_bytes) {
    _pDigest->update(pData, size_bytes);
}

/*!
 * \brief   Create and return a signature for a buffer.
 *
 * \param   pData       Data pointer.
 * \param   size_bytes   Size in bytes.
 *
 * \return  A base64 signature.
 */
std::string SignRsa::sign_data(const unsigned char* pData, std::size_t size_bytes) {
    _InternalDigest digest;
    // Guard not being setup or given bad things.
    if ((pData != nullptr) &&
        (size_bytes != 0)) {
        digest.update(pData, size_bytes);
    }
    return sign_digest(digest.get_digest());
}

/*!
 * \brief   Create and return a signature for a file.
 *
 * \param   pFile       File pointer.
 *
 * \return  A base64 signature.
 */
std::string SignRsa::sign_file(FILE* pFile) {
    // Guard not being setup or given bad things.
    if (pFile == nullptr) {
        return "";
    }

    // Stream file into buffer, update the digest.
    static unsigned char pBuffer[BUFFER_SIZE];
    _InternalDigest digest;
    std::size_t size_bytes = fread(pBuffer, 1, BUFFER_SIZE, pFile);
    std::size_t total = size_bytes;
    // Read data in chunks and feed it to OpenSSL SHA256
    while (size_bytes > 0)
    {
        digest.update(pBuffer, size_bytes);
        size_bytes = fread(pBuffer, 1, BUFFER_SIZE, pFile);
        total += size_bytes;
    }
    return sign_digest(digest.get_digest());
}

/*!
 * \brief   Create and return a signature for a file.
 *
 * \param   pDigest     Digest pointer.
 *
 * \return  A base64 signature.
 */
std::string SignRsa::sign_digest(const unsigned char* pDigest) {

    int result = 0; // Set to 0=invalid.

    std::string signature;

    // Guard not being setup or given bad things.
    // 
    // Invalid signature size, etc. are handled by RSA_verify, 
    // here we make sure our pointers are ok
    if ((_pDigest != nullptr) &&
        (_pPrivate != nullptr)) {

       char *pBuffer = (char *)malloc(RSA_size(_pPrivate));
       if (pBuffer != NULL) {

            unsigned int size_bytes = 0;
            
            // Ask the library to sign.
            result = RSA_sign(
                NID_sha256,
                pDigest,
                SHA256_DIGEST_LENGTH,
                (unsigned char *)pBuffer,
                &size_bytes,
                _pPrivate);

            if (result == 1) {
                signature = encodeBase64(pBuffer, size_bytes);                
            } else {
                signature.resize(0);
            }
        }
    }
    return signature;
}

/*!
 * \brief   Given buffer and signature, verify they match w/ public key.
 *
 * \param   pData       Data pointer.
 * \param   size_bytes  Size in bytes.
 * \param   signature   The signature.
 *
 * \return  True if they match, false if can't compare or failed.
 */
bool SignRsa::verify_data(const unsigned char* pData, std::size_t size_bytes, const std::string& signature) {

    _InternalDigest digest;

    // Guard not being setup or given bad things.
    if ((pData != nullptr) &&
        (size_bytes != 0)) {
        digest.update(pData, size_bytes);
    }
    return verify_digest(digest.get_digest(), signature);
}

/*!
 * \brief   Given a file and signature verify they match w/ public key.
 *
 * \param   pFile       File pointer.
 * \param   signature   The signature.
 *
 * \return  True if they match, false if can't compare or failed.
 */
bool SignRsa::verify_file(FILE* pFile, const std::string& signature) {
  
    // Guard not being setup or given bad things.
    if (pFile == nullptr) {
        return false;
    }
    
    // Stream file into buffer, update the digest.
    static unsigned char pBuffer[BUFFER_SIZE];
    _InternalDigest digest;
    std::size_t size_bytes = fread(pBuffer, 1, BUFFER_SIZE, pFile);

    // Read data in chunks and feed it to OpenSSL SHA256
    while (size_bytes > 0)
    {
        digest.update(pBuffer, size_bytes);
        size_bytes = fread(pBuffer, 1, BUFFER_SIZE, pFile);
    }

    // Verify the digest.
    return verify_digest(digest.get_digest(), signature);
}

/*!
 * \brief   Given a Digest and signature verify they match w/ public key.
 *
 * \param   pDigest     File pointer.
 * \param   signature   The signature.
 *
 * \return  True if they match, false if can't compare or failed.
 */
bool SignRsa::verify_digest(
    const unsigned char* pDigest,
    const std::string&  signature) {

    int result = 0; // Set to 0=invalid.

    // Guard not being setup or given bad things.
    // 
    // Invalid signature size, etc. are handled by RSA_verify, 
    // here we make sure our pointers are ok
    if ((_pDigest != nullptr) &&
        (_pPublic != nullptr)) {

        // Our signature is base64 without end of lines.

        // For readability, get the pointer and size here.
        const std::string decoded = decodeBase64(signature);
        const unsigned char * pSigBuffer = (const unsigned char* )decoded.data();
        const std::size_t  sig_size_bytes = decoded.size();
    
        // Ask the library to verify.
        result = RSA_verify(
            NID_sha256, 
            pDigest, 
            SHA256_DIGEST_LENGTH,
            pSigBuffer,
            (unsigned int)sig_size_bytes,
            _pPublic);
    }
    return (result == 1);
}

/*!
 * \brief   Finalize the context and create the digest bytes.
 *
 * \return  Digest bytes.
 */
const unsigned char* SignRsa::get_digest() const {
    return _pDigest->get_digest();
}

/*!
 * \brief   Decode base64 to data.
 * 
 * This is very fast and more permissive than OpenSSL's implimentation
 * and thereby makes it compatible with how Python does things with OpenSSL.
 * 
 * IF the input data is invalid, you get invalid output but it won't crash.
 * 
 * \param   base64_str      String.
 * 
 * \return  String.
 */
std::string SignRsa::decodeBase64(const std::string& base64_str) const {

    static constexpr char reverse_table[256] = {
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 62, 64, 64, 64, 63,
        52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 64, 64, 64, 64, 64, 64,

        64,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
        15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 64, 64, 64, 64, 64,
        64, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
        41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 64, 64, 64, 64, 64,

        // Unused unless a byte is invalid, we padd with
        // stuff we don't care so we can skip array indexing for speed.
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,

        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64,
        64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64
    };

    // This is meant to be simple and fast, but not validate the format.
    // In the RSA sense the rest of this code consumes the decoded 
    // data and validates it's content so ... format checking is not
    // useful it gets caught by the consumer.
    //
    // The fastest way to do this is process 3 input bytes at once
    // and construct the output with unrolling but when it's coded
    // that way it is 3-4x more complex with more corner cases to verify.
    // 
    // This is almost that fast (fast enough for the application) and
    // no corner cases.

    std::string     rv;         // Return value.
    unsigned int    acc = 0;    // Accumulate the value.
    unsigned int    bits = 0;   // Bits in the accumulated.
    for (unsigned char c : base64_str) {

        // This handles the problem that openSSL has, 
        // it either requires line breaks or forbids line breaks.
        if (::std::isspace(c) || c == '=') {
            continue;
        }

        // Decode the byte, accumulate 6 bits.
        char v = reverse_table[c];
        acc = (acc << 6) | v;
        bits += 6;

        // Export a byte if needed.
        if (bits >= 8) {
            bits -= 8;
            rv += (char)((acc >> bits) & 0xFF);
        }
    }

    // Return.
    return rv;
}

/*!
 * \brief   Encode base64 to data.
 *
 * IF the input data is invalid, you get invalid output but it won't crash.
 *
 * \param   raw_str     Raw data string.
 * \param   sizeBytes   Size in bytes.
 *
 * \return  String.
 */
std::string SignRsa::encodeBase64(const char* raw_str, std::size_t sizeBytes) const {
    static constexpr char encode_table[] = {
          'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
          'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
          'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
          'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
          'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
          'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
          'w', 'x', 'y', 'z', '0', '1', '2', '3',
          '4', '5', '6', '7', '8', '9', '+', '/'
    };

    std::string rv;

    std::size_t phase = 0;
    unsigned char lc = 0;
    unsigned char byte;
    for (std::size_t ii = 0; ii < sizeBytes; ii++) {
        // Every 3 characters = 4 bytes.    
        unsigned char c = (unsigned char) raw_str[ii];

        // NOTE: literals are used here instead of constexpr
        // because it is dramatically easier to understand.
        switch (phase) {
        case 0:
            // top 6 bits.
            byte = (c >> 2) & 0x3F;
            rv += encode_table[byte];
            break;
        case 1:
            // last 2 bits + 4 bits.
            byte = ((lc & 0x3) << 4) | (c >> 4);
            rv += encode_table[byte];
            break;
        default:
            // last 4 bits + 2 bits.
            byte = ((lc & 0xF) << 2) | (c >> 6);
            rv += encode_table[byte];

            // Last 6 bits.
            byte = (c & 0x3F);
            rv += encode_table[byte];
        }
        phase = (phase + 1) % 3;
        lc = c;
    }

    // Handle tail end.
    // 0 means we lined up luckily ideally, no padding.
    // 1 means we have 2 bits in lc to emit, two ==
    // 2 means we have 4 bits in lc to emit, one =
    if (phase == 1) {
        byte = ((lc & 0x3) << 4);
        rv += encode_table[byte];
        rv += "==";
    }
    if (phase == 2) {
        byte = (lc & 0xF) << 2;
        rv += encode_table[byte];
        rv += "=";
    }

    return rv;
}

/*!
 * \brief   Encode base64 to data.
 * 
 * IF the input data is invalid, you get invalid output but it won't crash.
 *
 * \param   raw_str     Raw data string.
 *
 * \return  String.
 */
std::string SignRsa::encodeBase64(const std::string& raw_str) const {
    return encodeBase64(raw_str.c_str(), raw_str.length());
}