#include "pch.h"
#include "BASE64.h"
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>


using namespace LenheartBase;


// 加解密函数签名
typedef int (*RSA_encryptOrDecrypt)(int flen, const unsigned char* from, unsigned char* to, RSA* rsa, int padding);


// 编解码转换表
unsigned char CBASE64::s_encTable[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
unsigned char CBASE64::s_decTable[] = {
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x3E, 0xFF, 0xFF, 0xFF, 0x3F,
    0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x3B, 0x3C, 0x3D, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E,
    0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28,
    0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 0x30, 0x31, 0x32, 0x33, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };

// 执行BASE64编码操作
std::string CBASE64::encode(const std::string& str)
{
    std::string strEncode;
    strEncode.resize((str.size() / 3 + 1) * 4);
    strEncode.resize(__encode((unsigned char*)strEncode.data(), (const unsigned char*)str.data(), str.size()));
    return strEncode;
}

// 执行BASE64解码操作
std::string CBASE64::decode(const std::string& str)
{
    std::string strDecode;
    strDecode.resize(str.size());
    strDecode.resize(__decode((unsigned char*)strDecode.data(), (const unsigned char*)str.data(), str.size()));
    return strDecode;
}

// 分组编码
int CBASE64::__encode(unsigned char* pDest, const unsigned char* pSrc, size_t nSrcLen)
{
    unsigned char* dst = pDest;
    const unsigned char* src = pSrc;
    size_t len = nSrcLen;

    unsigned char* odst = dst;
    unsigned long l = 0, m = 0, n = 0;

    // 循环处理分组
    size_t off = 0;
    for (; off + 3 <= len; off += 3)
    {
        l = *src++;
        m = *src++;
        n = *src++;

        *dst++ = s_encTable[(l >> 2) & 0x3F];
        *dst++ = s_encTable[((l << 4) & 0x30) | ((m >> 4) & 0x0F)];
        *dst++ = s_encTable[((m << 2) & 0x3C) | ((n >> 6) & 0x03)];
        *dst++ = s_encTable[n & 0x3F];
    }

    // 处理余下的2个字节
    if (off + 2 <= len)
    {
        l = *src++;
        m = *src++;

        *dst++ = s_encTable[(l >> 2) & 0x3F];
        *dst++ = s_encTable[((l << 4) & 0x30) | ((m >> 4) & 0x0F)];
        *dst++ = s_encTable[((m << 2) & 0x3C)];
        *dst++ = '=';
    }

    // 处理余下的1个字节
    else if (off + 1 <= len)
    {
        l = *src++;

        *dst++ = s_encTable[(l >> 2) & 0x3F];
        *dst++ = s_encTable[((l << 4) & 0x30)];
        *dst++ = '=';
        *dst++ = '=';
    }

    return (int)(dst - odst);
}

// 分组解码
int CBASE64::__decode(unsigned char* pDest, const unsigned char* pSrc, size_t nSrcLen)
{
    unsigned char* dst = pDest;
    const unsigned char* src = pSrc;
    size_t len = nSrcLen;

    unsigned char* odst = dst;
    unsigned long l = 0, m = 0, n = 0, o = 0;

    // 循环处理分组
    size_t off = 0;
    for (; off + 4 <= len; off += 4)
    {
        if ((src[0] > 0x7F) || (src[1] > 0x7F) || (src[2] > 0x7F) || (src[3] > 0x7F))
            return (int)(dst - odst);
        if ((src[0] == '=') || (src[1] == '=') || (src[2] == '=') || (src[3] == '='))
            break;

        l = s_decTable[*src++];
        m = s_decTable[*src++];
        n = s_decTable[*src++];
        o = s_decTable[*src++];

        *dst++ = (unsigned char)(((l << 2) & 0xFC) | ((m >> 4) & 0x03));
        *dst++ = (unsigned char)(((m << 4) & 0xF0) | ((n >> 2) & 0x0F));
        *dst++ = (unsigned char)(((n << 6) & 0xC0) | (o & 0x3F));
    }

    // 处理余下的3个字节
    if (off + 3 <= len)
    {
        if ((src[0] != '=') && (src[1] != '='))
        {
            l = s_decTable[*src++];
            m = s_decTable[*src++];

            *dst++ = (unsigned char)(((l << 2) & 0xFC) | ((m >> 4) & 0x03));
        }

        if ((src[2] != '='))
        {
            n = s_decTable[*src++];

            *dst++ = (unsigned char)(((m << 4) & 0xF0) | ((n >> 2) & 0x0F));
        }
    }

    // 处理余下的两个字节
    else if (off + 2 <= len)
    {
        if ((src[0] != '=') && (src[1] != '='))
        {
            l = s_decTable[*src++];
            m = s_decTable[*src++];

            *dst++ = (unsigned char)(((l << 2) & 0xFC) | ((m >> 4) & 0x03));
        }
    }

    return (int)(dst - odst);
}

std::string CBASE64::RsaPriDecrypt(const std::string& cipher_text, const std::string& pri_key)
{
    std::string decrypt_text;
    RSA* rsa = RSA_new();
    BIO* keybio;
    keybio = BIO_new_mem_buf((unsigned char*)pri_key.c_str(), -1);

    rsa = PEM_read_bio_RSAPrivateKey(keybio, &rsa, NULL, NULL);
    if (rsa == nullptr)
    {
        unsigned long err = ERR_get_error(); //获取错误号
        char err_msg[1024] = { 0 };
        ERR_error_string(err, err_msg); // 格式：error:errId:库:函数:原因
        return std::string();
    }

    // 获取RSA单次处理的最大长度
    int key_len = RSA_size(rsa);
    char* sub_text = new char[key_len + 1];
    memset(sub_text, 0, key_len + 1);
    int ret = 0;
    std::string sub_str;
    unsigned int pos = 0;
    // 对密文进行分段解密
    while (pos < cipher_text.length() - 1)
    {
        sub_str = cipher_text.substr(pos, key_len);
        memset(sub_text, 0, key_len + 1);
        ret = RSA_private_decrypt(sub_str.length(), (const unsigned char*)sub_str.c_str(), (unsigned char*)sub_text, rsa, 1);
        if (ret >= 0)
        {
            decrypt_text.append(std::string(sub_text, ret));
            // printf("pos:%d, Length: %d ,sub: %s\n", pos, cipher_text.length(),sub_text);
            pos += key_len;
        }
    }
    // 释放内存
    delete[] sub_text;
    BIO_free_all(keybio);
    RSA_free(rsa);

    return decrypt_text;
}

std::string CBASE64::RsaPriEncrypt(const std::string& clear_text, const std::string& pri_key)
{
    std::string encrypt_text;
    BIO* keybio = BIO_new_mem_buf((unsigned char*)pri_key.c_str(), -1);
    RSA* rsa = RSA_new();
    rsa = PEM_read_bio_RSAPrivateKey(keybio, &rsa, NULL, NULL);
    if (!rsa)
    {
        BIO_free_all(keybio);
        return std::string("");
    }

    // 获取RSA单次可以处理的数据块的最大长度
    int key_len = RSA_size(rsa);
    int block_len = key_len - 11; // 因为填充方式为RSA_PKCS1_PADDING, 所以要在key_len基础上减去11

    // 申请内存：存贮加密后的密文数据
    char* sub_text = new char[key_len + 1];
    memset(sub_text, 0, key_len + 1);
    int ret = 0;
    unsigned int pos = 0;
    std::string sub_str;
    // 对数据进行分段加密（返回值是加密后数据的长度）
    while (pos < clear_text.length())
    {
        sub_str = clear_text.substr(pos, block_len);
        memset(sub_text, 0, key_len + 1);
        ret = RSA_private_encrypt(sub_str.length(), (const unsigned char*)sub_str.c_str(), (unsigned char*)sub_text, rsa, RSA_PKCS1_PADDING);
        if (ret >= 0)
        {
            encrypt_text.append(std::string(sub_text, ret));
        }
        pos += block_len;
    }

    // 释放内存
    delete sub_text;
    BIO_free_all(keybio);
    RSA_free(rsa);

    return encrypt_text;
}

// 生成RSA密钥结构
void* CBASE64::__genKey(int nBits)
{
    RSA* rsa = RSA_new();

    BIGNUM* bne = BN_new();
    BN_set_word(bne, RSA_F4);

    int ret = RSA_generate_key_ex(rsa, nBits, bne, NULL);

    BN_free(bne);

    if (ret != 1)
    {
        RSA_free(rsa);
        return NULL;
    }

    return (void*)rsa;
}

// 生成密钥对，输出为PEM字符串
bool CBASE64::genKeyStrings(std::string& strPrivateKeyPEMString, std::string& strPublicKeyPEMString)
{
    // 生成RSA
    RSA* rsa = (RSA*)__genKey(1024);
    if (rsa == NULL)
        return false;

    // 输出私钥
    {
        BIO* bio = BIO_new(BIO_s_mem());
        int ret = PEM_write_bio_RSAPrivateKey(bio, rsa, NULL, NULL, 0, NULL, NULL);
        if (ret != 1)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        char sBuf[1024] = { 0 };
        int bytes = BIO_read(bio, sBuf, 1024);
        if (bytes <= 0)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        BIO_free(bio);
        strPrivateKeyPEMString.assign(sBuf, bytes);
    }

    // 输出公钥
    {
        BIO* bio = BIO_new(BIO_s_mem());
        int ret = PEM_write_bio_RSAPublicKey(bio, rsa);
        if (ret != 1)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        char sBuf[1024] = { 0 };
        int bytes = BIO_read(bio, sBuf, 1024);
        if (bytes <= 0)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        BIO_free(bio);
        strPublicKeyPEMString.assign(sBuf, bytes);
    }

    RSA_free(rsa);
    return true;
}


// 使用PEM私钥字符串加密
bool CBASE64::encryptByPrivatePEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载私钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPrivateKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 使用私钥加密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_private_encrypt, true);
    RSA_free(rsa);
    return b;
}


// 使用RSA执行加密或解密
bool CBASE64::__encryptOrDecrypt(const std::string& strIn, std::string& strOut, const void* pRSA, const void* pFunc, bool bEncrypt)
{
    const RSA_encryptOrDecrypt encryptOrDecrypt = (const RSA_encryptOrDecrypt)pFunc;

    // 计算加密块大小
    int nKeySize = RSA_size((RSA*)pRSA);
    int nBlockSize = bEncrypt ? (nKeySize - RSA_PKCS1_PADDING_SIZE) : nKeySize;

    const unsigned char* pIn = (const unsigned char*)strIn.data();
    int nInSize = strIn.size();

    unsigned char* pBuf = new unsigned char[nKeySize];

    // 分组迭代加密
    for (int i = 0; i < nInSize; )
    {
        int nBlockLen = (nInSize - i > nBlockSize ? nBlockSize : nInSize - i);

        int ret = encryptOrDecrypt(nBlockLen, pIn + i, pBuf, (RSA*)pRSA, RSA_PKCS1_PADDING);
        if (ret <= 0)
        {
            delete[] pBuf;
            return false;
        }

        strOut.append((char*)pBuf, ret);
        i += nBlockLen;
    }

    delete[] pBuf;
    return true;
}

// 使用PEM公钥字符串加密
bool CBASE64::encryptByPublicPEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载公钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 使用公钥加密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_public_encrypt, true);
    RSA_free(rsa);
    return b;
}

// 使用PEM公钥字符串解密
bool CBASE64::decryptByPublicPEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载公钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 检查密文大小是否为分组的整数倍
    int keySize = RSA_size(rsa);
    int blockSize = keySize;
    if ((strIn.size() % blockSize) != 0)
        return false;

    // 使用公钥解密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_public_decrypt, false);
    RSA_free(rsa);
    return b;
}