#include "FlashFileUtils.h"
#include "openssl/rsa.h"
#include "openssl/pem.h"
#include "openssl/err.h"
#include "openssl/aes.h"
#include "openssl/rsa.h"
#include "openssl/err.h"
#include "openssl/x509.h"
#include <iostream>
#include <fstream>
#include <cstring>
#include "MyStringUtils.h"
#include "algorithm.h"

#define modp_b64_encode_len(A) ((A+2)/3 * 4 + 1)
#define modp_b64_decode_len(A) (A / 4 * 3 + 2)


CFlashFileUtils::CFlashFileUtils()
{
	m_strInitVector = "64656275673030303030303030303030";
	m_bInitRSA = false;
	memset(m_Key, 0, 1024);
}

CFlashFileUtils::~CFlashFileUtils()
{
}

bool CFlashFileUtils::InitRSA(const string& strRSAPrivateKeyFileName, const string& strKeyInfo)
{
	bool bRet = false;
	m_bInitRSA = false;
	memset(m_Key, 0, 1024);

	if (strKeyInfo.empty())
	{
		bRet = false;
		return bRet;
	}

	bRet = RSA_Decrypt(m_Key, strRSAPrivateKeyFileName, strKeyInfo);
	m_bInitRSA = bRet;

	return bRet;
}

bool CFlashFileUtils::DecryptFlashFile(unsigned char* pOutBuffer, unsigned char* pDataBuffer, unsigned int& nDataLen)
{
	bool bRet = false;

	if (!m_bInitRSA)
	{
		return bRet;
	}

	bRet = AES_DecryptEx(pOutBuffer, pDataBuffer, nDataLen);
	if (!bRet)
	{
		return bRet;
	}

	//Dump("AES_DecryptEx.bin", pOutBuffer, nDataLen);

	nDataLen =DePKCS7Padding(pOutBuffer, nDataLen);

	//Dump("AES_DecryptEx_VBF.bin", pOutBuffer, nDataLen);

	return bRet;
}

bool CFlashFileUtils::EncryptRSAPrivateKey(unsigned char* pOutBuffer, unsigned int& nOutLen, unsigned char* pDataBuffer, unsigned int nDataLen)
{
	bool bRet = false;

	nDataLen += PKCS7Padding(pOutBuffer, nDataLen);

	bRet = AES_EncryptRSAPrivateKey(pOutBuffer, pDataBuffer, nDataLen);
	if (!bRet)
	{
		return bRet;
	}

	nOutLen = nDataLen;

	//Dump("DiagnosticData.bin", pOutBuffer, nDataLen);

	return bRet;
}

bool CFlashFileUtils::DecryptRSAPrivateKey(unsigned char* pOutBuffer, unsigned char* pDataBuffer, unsigned int& nDataLen)
{
	bool bRet = false;

	bRet = AES_DecryptRSAPrivateKey(pOutBuffer, pDataBuffer, nDataLen);
	if (!bRet)
	{
		return bRet;
	}

	//Dump("AES_DecryptEx.bin", pOutBuffer, nDataLen);

	nDataLen = DePKCS7Padding(pOutBuffer, nDataLen);

	//Dump("DePKCS7Padding.bin", pOutBuffer, nDataLen);

	return bRet;
}

bool CFlashFileUtils::DecryptConfig(string& strConfig, const string& strBuffer)
{
	bool bRet = false;

	//njb 20230513 add 关键代码加壳
#ifdef _VM_PROTECT_
	VMStart();
#endif // _WINDOWS_
	int nBufferLen = (int)strBuffer.size();
	size_t dstlen = modp_b64_decode_len(nBufferLen);
	unsigned char* source = (unsigned char*)malloc(dstlen);
	if (source == NULL) {
		return false;
	}
	memset(source, 0, dstlen);
	nBufferLen = base64_decode((const char*)strBuffer.c_str(), (long)strBuffer.size(), (unsigned char*)source, (int)dstlen);
	bRet = AES_DecryptConfig(source, source, nBufferLen);
	if (!bRet)
	{
		free(source);
		source = NULL;

		return bRet;
	}

	nBufferLen = DePKCS7Padding(source, nBufferLen);

	strConfig.assign((char*)source, nBufferLen);

	//Dump("Config.bin", source, nBufferLen);

	free(source);
	source = NULL;

	//njb 20230513 add 关键代码加壳
#ifdef _VM_PROTECT_
	VMEnd();
#endif // _WINDOWS_
	return bRet;
}

bool CFlashFileUtils::LoadConfig(std::string& strConfig, const std::string& strFileName)
{
	bool bRet = false;

	int nDataLen = 0;
	if (strFileName.empty())
	{
		bRet = false;
		return bRet;
	}

	FILE* fp = NULL;
	fp = fopen(strFileName.c_str(), "rb");

	if (fp == NULL)
	{
		bRet = false;
		return bRet;
	}

	fseek(fp, 0, SEEK_END);
	long datasize = ftell(fp);
	fseek(fp, 0, SEEK_SET);

	unsigned char* pBuffer = new unsigned char[datasize + 1];
	memset(pBuffer, 0, datasize + 1);

	fread(pBuffer, 1, datasize, fp);
	fclose(fp);

	std::string strBuffer;
	strBuffer.assign((char*)pBuffer, datasize);

	delete[] pBuffer;
	pBuffer = NULL;

	bRet = DecryptConfig(strConfig, strBuffer);
	return bRet;
}

static RSA* Buffer2EVPkey(const uint8_t* key, size_t key_len)
{
	BIO* b = BIO_new_mem_buf((void*)key, (int)key_len);
	if (NULL == b) {
		return NULL;
	}

	//如果私钥文件使用了密码进行了加密，则下面函数的第二个和第三个入参要赋值，具体请参见openssl的PEM_read_bio_PrivateKey接口描述
	EVP_PKEY* evpkey = PEM_read_bio_PrivateKey(b, NULL, NULL, NULL);
	if (NULL == evpkey) {
		BIO_free(b);
		return NULL;
	}
	BIO_free(b);

	RSA* rsa = EVP_PKEY_get1_RSA(evpkey);
	EVP_PKEY_free(evpkey);

	return rsa;
}

bool CFlashFileUtils::RSA_Decrypt(unsigned char* key, const string& strRSAPrivateKeyFileName, const std::string& strKeyInfo)
{
	size_t nKeyLen = strKeyInfo.size();

	size_t dstlen = modp_b64_decode_len(nKeyLen);
	unsigned char* source = (unsigned char*)malloc(dstlen);
	memset(source, 0, dstlen);
	nKeyLen = base64_decode((const char*)strKeyInfo.c_str(), (long)strKeyInfo.size(), (unsigned char*)source, (int)nKeyLen);

	RSA* pRSA;
	int nRSA_len;
	//FILE* file;

	{
		FILE* fp = nullptr;
		fp = fopen(strRSAPrivateKeyFileName.c_str(), "rb");
		if (fp == nullptr)
		{
			free(source);
			source = NULL;

			return false;
		}

		fseek(fp, 0, SEEK_END);
		unsigned int datasize = ftell(fp);
		fseek(fp, 0, SEEK_SET);

		unsigned char* pBuffer = new unsigned char[datasize + 1];
		memset(pBuffer, 0, datasize + 1);

		fread(pBuffer, 1, datasize, fp);
		fclose(fp);

		DecryptRSAPrivateKey(pBuffer, pBuffer, datasize);

		pRSA = Buffer2EVPkey(pBuffer, datasize);

		delete[] pBuffer;
		pBuffer = NULL;

		if (pRSA == NULL)
		{
			free(source);
			source = NULL;

			return false;
		}
	}

	nRSA_len = RSA_size(pRSA);

	if (nRSA_len > 1024)
	{
		free(source);
		source = NULL;

		RSA_free(pRSA);
		//fclose(file);
		return false;
	}

	memset(key, 0, nRSA_len + 1);

	if (RSA_private_decrypt(nRSA_len, (unsigned char*)source, (unsigned char*)key, pRSA, RSA_PKCS1_PADDING) < 0)
	{
		if (source != NULL)
		{
			free(source);
			source = NULL;
		}

		RSA_free(pRSA);
		//fclose(file);
		return false;
	}

	RSA_free(pRSA);

	if (source != NULL)
	{
		free(source);
		source = NULL;
	}

	//fclose(file);

	return true;
}

bool CFlashFileUtils::AES_EncryptRSAPrivateKey(unsigned char* pOutBuffer, unsigned char* pDataBuffer, int nDataLen)
{
	bool bRet = false;
	unsigned char key[AES_BLOCK_SIZE] = { 0xAE, 0x15, 0x62, 0x75, 0x67, 0xBC, 0xA2, 0x98, 0x23, 0x11, 0xDC, 0xF1, 0xFA, 0xAA, 0xBC, 0x54 };        // AES_BLOCK_SIZE = 16
	unsigned char iv[AES_BLOCK_SIZE] = { 0x64, 0x65, 0x62, 0x75, 0x67, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30 };        // init vector

	AES_KEY aes;

	if (AES_set_encrypt_key(key, 128, &aes) < 0)
	{
		bRet = false;
		return bRet;
	}

	AES_cbc_encrypt(pDataBuffer, pOutBuffer, nDataLen, &aes, iv, AES_ENCRYPT);

	bRet = true;
	return bRet;
}

bool CFlashFileUtils::AES_DecryptRSAPrivateKey(unsigned char* pOutBuffer, unsigned char* pDataBuffer, int nDataLen)
{
	bool bRet = false;
	unsigned char key[AES_BLOCK_SIZE] = { 0xAE, 0x15, 0x62, 0x75, 0x67, 0xBC, 0xA2, 0x98, 0x23, 0x11, 0xDC, 0xF1, 0xFA, 0xAA, 0xBC, 0x54 };        // AES_BLOCK_SIZE = 16
	unsigned char iv[AES_BLOCK_SIZE] = { 0x64, 0x65, 0x62, 0x75, 0x67, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30 };        // init vector

	AES_KEY aes;

	if (AES_set_decrypt_key(key, 128, &aes) < 0)
	{
		bRet = false;
		return bRet;
	}

	AES_cbc_encrypt(pDataBuffer, pOutBuffer, nDataLen, &aes, iv, AES_DECRYPT);

	bRet = true;
	return bRet;
}

bool CFlashFileUtils::AES_DecryptConfig(unsigned char* pOutBuffer, unsigned char* pDataBuffer, int nDataLen)
{
	bool bRet = false;
	unsigned char key[AES_BLOCK_SIZE] = { 'j', 'i', 'd', 'u', 'a', 'u', 't', 'o', 'j', 'i', 'd', 'u', 'a', 'u', 't', 'o' };        // AES_BLOCK_SIZE = 16
	unsigned char iv[AES_BLOCK_SIZE] = { 0x00 };        // init vector

	AES_KEY aes;

	if (AES_set_decrypt_key(key, 128, &aes) < 0)
	{
		bRet = false;
		return bRet;
	}

	AES_cbc_encrypt(pDataBuffer, pOutBuffer, nDataLen, &aes, iv, AES_DECRYPT);

	bRet = true;
	return bRet;
}

bool CFlashFileUtils::AES_DecryptEx(unsigned char* pOutBuffer, unsigned char* pDataBuffer, int nDataLen)
{
	bool bRet = false;
	unsigned char key[AES_BLOCK_SIZE];        // AES_BLOCK_SIZE = 16
	unsigned char iv[AES_BLOCK_SIZE] = { 0x64, 0x65, 0x62, 0x75, 0x67, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30 };        // init vector

	AES_KEY aes;

	// Generate AES 128-bit key
	for (int i = 0; i < 16; ++i)
	{
		key[i] = m_Key[i];
	}

	if (AES_set_decrypt_key(key, 128, &aes) < 0)
	{
		bRet = false;
		return bRet;
	}

	AES_cbc_encrypt(pDataBuffer, pOutBuffer, nDataLen, &aes, iv, AES_DECRYPT);

	bRet = true;
	return bRet;
}

int CFlashFileUtils::PKCS7Padding(unsigned char* pBuffer, int nDataLen)
{
	int remain = nDataLen % 16;

	if (remain > 0)
	{
		remain = 16 - remain;

		for (int i = 0; i < remain; i++)
		{
			pBuffer[nDataLen + i] = remain;
		}
	}

	return remain;
}

int CFlashFileUtils::DePKCS7Padding(unsigned char* pBuffer, int nDataLen)
{
	if (pBuffer == NULL || nDataLen == 0)
	{
		return 0;
	}

	int remain = pBuffer[nDataLen - 1];
	if (remain > nDataLen || remain > 0x10)
	{
		return 0;
	}

	return nDataLen - remain;
}

void CFlashFileUtils::Dump(std::string strFileName, unsigned char* pBuffer, int nDataLen)
{
	FILE* fp = NULL;
	fp = fopen(strFileName.c_str(), "wb");
	fwrite(pBuffer, nDataLen, 1, fp);
	fclose(fp);
}