// RVCComm.cpp : Defines the entry point for the DLL application.
//

#include "stdafx.h"
#include "RVCComm.h"
#include "CMBSMDLL.h"

#include "ClientComm.h"

#if defined(_MSC_VER)
#include <winsock2.h>
#endif //_MSC_VER

#include <assert.h>
#include "rsa.h"
#include "md5.h"
#include "pem.h"
#include <err.h>
#include "dbgutil.h"

#if defined(_MSC_VER)
#pragma comment(lib, "Ws2_32.lib")
#else
#include <winpr/locale.h>
#include "utils.h"
#endif //_MSC_VER

#define TAG RVCCOMM_TAG("rvccomm")


CSecureClientBase::CSecureClientBase()
{
	m_nRefCount = 1;		// Entity持有, CClientComm生命期由CSecureClientBase控制，故持有不占用计数
	m_pClientComm = new CClientComm(this);
}

bool InitSocket()
{
	//Initialize Winsock
	WSADATA wsaData;

	int nResult = WSAStartup(MAKEWORD(2,2), &wsaData);
	return nResult ==0;
}


void UninitSocket()
{
	WSACleanup();
}
#ifdef _WIN32
BOOL APIENTRY DllMain( HANDLE hModule, 
                       DWORD  ul_reason_for_call, 
                       LPVOID lpReserved
					 )
{
	bool bRet(true);

    switch (ul_reason_for_call)
	{
		case DLL_PROCESS_ATTACH:
			OpenSSL_add_all_ciphers();
			OpenSSL_add_all_digests();
			ERR_load_crypto_strings();

			bRet = InitSocket();
			break;

		case DLL_THREAD_ATTACH:
			break;

		case DLL_THREAD_DETACH:
			break;

		case DLL_PROCESS_DETACH:
			UninitSocket();
			break;
    }

    return bRet ? TRUE : FALSE;
}
#endif


// 客户端通过此函数获得通讯类实例
/*
RVCCOMM_API IClientCommFunc* GetClientComm(CSecureClientBase *pCallback)
{
	TOOLKIT_ASSERT(pCallback != NULL);
	if (pCallback == NULL)
		return NULL;

	pCallback->m_pClientComm = new CClientComm(pCallback);
	return pCallback->m_pClientComm;
}

// 服务端通过此函数获得通讯类实例
RVCCOMM_API CServerCommBase* GetServerComm(CServerCallback *pCallback)
{
	TOOLKIT_ASSERT(pCallback != NULL);
	if (pCallback == NULL)
		return NULL;

	pCallback->m_pServerComm = new CServerComm(pCallback);
	return pCallback->m_pServerComm;
}*/


string GetSysErrorMsg(int nErrorCode)
{
	char pMsg[1024];
	memset(pMsg, 0, sizeof(pMsg));

	if(::FormatMessage(
			FORMAT_MESSAGE_IGNORE_INSERTS | FORMAT_MESSAGE_FROM_SYSTEM,
			NULL,
			nErrorCode,
			MAKELANGID(LANG_NEUTRAL,SUBLANG_DEFAULT),
			pMsg,
			sizeof(pMsg)/sizeof(char),
			NULL))
	{
		pMsg[sizeof(pMsg)/sizeof(char)-1]=0;

		int nLen = strlen(pMsg);
		if (pMsg[nLen-1] == '\n' && pMsg[nLen-2] == '\r')
			pMsg[nLen-2] = 0;
	}

	return pMsg;
}


bool CreateRsaKeyPair(BYTE *pPubKeyBuf, int *pPubKeyBufLen, BYTE *pPriKeyBuf, int *pPriKeyBufLen)
{
	if (*pPubKeyBufLen < 140 || *pPriKeyBufLen <608)
		return false;

	memset(pPubKeyBuf, 0, *pPubKeyBufLen);
	memset(pPriKeyBuf, 0, *pPriKeyBufLen);

	RSA *rsa = RSA_generate_key(1024, RSA_F4, NULL, NULL);
	if (rsa == NULL)
		return false;

	*pPubKeyBufLen = i2d_RSAPublicKey(rsa, &pPubKeyBuf);
	*pPriKeyBufLen = i2d_RSAPrivateKey(rsa, &pPriKeyBuf);

	RSA_free(rsa);
	return (*pPubKeyBufLen >0 && *pPriKeyBufLen >0);
}

// 使用RSA公钥加密
bool EncWithRsaPubKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, RSA* rsa_pub)
{
	// 公钥加密
	TOOLKIT_ASSERT(*pDestLen >= nSourceLen);

	BYTE *pFrom = pSource;
	BYTE *pEnd = pSource + nSourceLen;
	int nTotalLen = 0;

	while(pFrom < pEnd)
	{
		if (*pDestLen - nTotalLen < 128)
		{
			RSA_free(rsa_pub);
			return false;
		}
			
		int nBlockLen = (pEnd - pFrom) <= 117 ? (pEnd - pFrom) : 117;
		int nRet = RSA_public_encrypt(nBlockLen, pFrom, pDest + nTotalLen, rsa_pub, RSA_PKCS1_PADDING);
		if (nRet <= 0)
		{
			RSA_free(rsa_pub);
			return false;
		}

		pFrom += nBlockLen;
		nTotalLen += nRet;
	}

	*pDestLen = nTotalLen;
	RSA_free(rsa_pub);
	return true;
}

// 使用RSA公钥加密
bool EncWithRsaPubKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, BYTE *pRsaPubKey, int nKeyLen)
{
	// 还原RSA公钥
	const BYTE *p = pRsaPubKey;
	RSA *rsa_pub = d2i_RSAPublicKey(NULL, &p, nKeyLen);

	if (rsa_pub == NULL)
		return false;

	return EncWithRsaPubKey(pSource, nSourceLen, pDest, pDestLen, rsa_pub);
}

// 使用RSA私钥加密
bool EncWithRsaPriKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, RSA *rsa_pri)
{
	// 私钥加密
	TOOLKIT_ASSERT(*pDestLen >= nSourceLen);

	BYTE *pFrom = pSource;
	BYTE *pEnd = pSource + nSourceLen;
	int nTotalLen = 0;

	while(pFrom < pEnd)
	{
		if (*pDestLen - nTotalLen < 128)
		{
			RSA_free(rsa_pri);
			return false;
		}
			
		int nBlockLen = (pEnd - pFrom) <= 117 ? (pEnd - pFrom) : 117;
		int nRet = RSA_private_encrypt(nBlockLen, pFrom, pDest + nTotalLen, rsa_pri, RSA_PKCS1_PADDING);
		if (nRet <= 0)
		{
			RSA_free(rsa_pri);
			return false;
		}

		pFrom += nBlockLen;
		nTotalLen += nRet;
	}

	*pDestLen = nTotalLen;
	RSA_free(rsa_pri);
	return true;
}

// 使用RSA私钥加密
bool EncWithRsaPriKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, BYTE *pRsaPriKey, int nKeyLen)
{
	// 还原RSA私钥
	const BYTE *p = pRsaPriKey;
	RSA *rsa_pri = d2i_RSAPrivateKey(NULL, &p, nKeyLen);

	if (rsa_pri == NULL)
		return false;

	return EncWithRsaPriKey(pSource, nSourceLen, pDest, pDestLen, rsa_pri);
}

// 使用RSA公钥解密
bool DecWithRsaPubKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, RSA *rsa_pub)
{
	// 公钥解密
	TOOLKIT_ASSERT(*pDestLen >= nSourceLen);
	TOOLKIT_ASSERT(nSourceLen % 128 ==0);

	BYTE *pFrom = pSource;
	BYTE *pEnd = pSource + nSourceLen;
	int nTotalLen = 0;

	while(pFrom < pEnd)
	{
		if (*pDestLen - nTotalLen < 128)
		{
			RSA_free(rsa_pub);
			return false;
		}
			
		int nRet = RSA_public_decrypt(128, pFrom, pDest + nTotalLen, rsa_pub, RSA_PKCS1_PADDING);
		if (nRet <= 0)
		{
			RSA_free(rsa_pub);
			return false;
		}

		pFrom += 128;
		nTotalLen += nRet;
	}

	*pDestLen = nTotalLen;
	RSA_free(rsa_pub);
	return true;
}

// 使用RSA公钥解密
bool DecWithRsaPubKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, BYTE *pRsaPubKey, int nKeyLen)
{
	// 还原RSA公钥
	const BYTE *p = pRsaPubKey;
	RSA *rsa_pub = d2i_RSAPublicKey(NULL, &p, nKeyLen);

	if (rsa_pub == NULL)
		return false;

	return DecWithRsaPubKey(pSource, nSourceLen, pDest, pDestLen, rsa_pub);
}

// 使用RSA私钥解密
bool DecWithRsaPriKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, RSA *rsa_pri)
{
	// 私钥解密
	TOOLKIT_ASSERT(*pDestLen >= nSourceLen);
	TOOLKIT_ASSERT(nSourceLen % 128 ==0);

	BYTE *pFrom = pSource;
	BYTE *pEnd = pSource + nSourceLen;
	int nTotalLen = 0;

	while(pFrom < pEnd)
	{
		if (*pDestLen - nTotalLen < 128)
		{
			RSA_free(rsa_pri);
			return false;
		}
			
		int nRet = RSA_private_decrypt(128, pFrom, pDest + nTotalLen, rsa_pri, RSA_PKCS1_PADDING);
		if (nRet <= 0)
		{
			RSA_free(rsa_pri);
			return false;
		}

		pFrom += 128;
		nTotalLen += nRet;
	}

	*pDestLen = nTotalLen;
	RSA_free(rsa_pri);
	return true;
}

// 使用RSA私钥解密
bool DecWithRsaPriKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen, BYTE *pRsaPriKey, int nKeyLen)
{
	// 还原RSA私钥
	const BYTE *p = pRsaPriKey;
	RSA *rsa_pri = d2i_RSAPrivateKey(NULL, &p, nKeyLen);

	if (rsa_pri == NULL)
		return false;

	return DecWithRsaPriKey(pSource, nSourceLen, pDest, pDestLen, rsa_pri);
}


bool MD5Hash(BYTE *pData, int nLen, BYTE md5[16])
{
	MD5_CTX context;
    memset(md5, 0, 16);
    
    MD5_Init(&context);
    MD5_Update(&context, pData, nLen);
    MD5_Final(md5, &context);
	return true;
}

///*TODO(80374374@3/22/2023):  Bcdhex_2_aschex 的移除，涉及到 util.cpp 文件的删除*/
bool SM3Hash(BYTE* pData, int nLen, BYTE hash[32])
{
#if defined(_MSC_VER)
	return 0 == CMBSM3Digest((unsigned char*)pData, nLen, hash);
#else
	WLog_DBG(TAG, "Enter %s", __FUNCTION__);
	unsigned char msg[1024] = { 0 };
	//Bcdhex_2_aschex(pData, nLen, msg);
	//WLog_DBG(TAG, "SM3Plain:%s", msg);
	WLog_DBG(TAG, "to CMBSM3Digest, msglen: %d, %d, %d", nLen, sizeof(unsigned char), sizeof(BYTE));
	unsigned char outmsgbcd[33] = { 0 };
	int ret = CMBSM3Digest((unsigned char*)pData, nLen, outmsgbcd);
	if (ret != 0) {
		WLog_ERR(TAG, "CMBSM3Digest failed: %d", ret);
		return false;
	} else {
		unsigned char outmsg[65] = { 0 };
		Bcdhex_2_aschex(outmsgbcd, 32, outmsg);
		WLog_DBG(TAG, "SM3:%s", outmsg);
		memcpy(hash, outmsgbcd, sizeof(BYTE) * 32);
		return true;
	}
#endif //_MSC_VER
}

char* GetSMVersion()
{
	return CMBSMGetVersion();
}

bool CreateSM2KeyPair(BYTE* pPubKeyBuf, int* pPubKeyBufLen, BYTE* pPriKeyBuf, int* pPriKeyBufLen)
{
	return 0 == CMBSM2KeyGen(pPubKeyBuf, pPubKeyBufLen, pPriKeyBuf, pPriKeyBufLen);
}

bool EncWithSM2PubKey(BYTE* pSource, int nSourceLen, BYTE* pDest, int* pDestLen, BYTE* pPubKey, int nKeyLen)
{
	return 0 == CMBSM2Encrypt(pPubKey, nKeyLen, pSource, nSourceLen, pDest, pDestLen);
}

bool DecWithSM2PriKey(BYTE* pSource, int nSourceLen, BYTE* pDest, int* pDestLen, BYTE* pPriKey, int nKeyLen)
{
	return 0 == CMBSM2Decrypt(pPriKey, nKeyLen, pSource, nSourceLen, pDest, pDestLen);
}

bool  SM2SignWithSM3(unsigned char* privkey, int privkey_len, unsigned char* msg, int msg_len, unsigned char* signature, int* sig_len)
{
	return 0 == CMBSM2SignWithSM3(privkey, privkey_len, msg, msg_len, signature, sig_len);
}

bool  SM2VerifyWithSM3(unsigned char* pubkey, int pubkey_len, unsigned char* msg, int msg_len, unsigned char* signature, int sig_len)
{
	return 0 == CMBSM2VerifyWithSM3(pubkey, pubkey_len, msg, msg_len, signature, sig_len);
}


void GenerateSM4_ECBkey(std::string keyStr, BYTE key[16])
{
	BYTE tempKey[32] = { 0 };
	if (keyStr.length() > 0)
	{
		SM3Hash((BYTE*)keyStr.c_str(), keyStr.length() > 32 ? 32 : keyStr.length(), tempKey);
		memcpy(key, tempKey, 16);
	}
	else
		memset(key, 0, 16);
	
	//Aschex_2_bcdhex(tempKey, 32, key);
}

bool subByteArr(BYTE* src, BYTE* dst, int offset, int size)
{
	if (nullptr == src || nullptr == dst)
		return false;

	memcpy(dst, src + offset, size);
	return true;
}


RVCCOMM_API bool EncWithSM4_ECB(std::string keyStr, BYTE* input, int length, BYTE* output, int* output_len)
{
	unsigned char key[17] = { 0 };
	GenerateSM4_ECBkey(keyStr, key);
	//printfHEX(key, 16);
	return 0 == CMBSM4EncryptWithECB(key, input, length, output, output_len);
	
}

bool DecWithSM4_ECB(std::string keyStr, BYTE* input, int length, BYTE* output, int* output_len)
{
	unsigned char key[17] = { 0 };
	GenerateSM4_ECBkey(keyStr, key);
	return 0 == CMBSM4DecryptWithECB(key, input, length, output, output_len);
}


bool EncWithSM4_ECB(BYTE key[16], BYTE* input, int length, BYTE* output, int* output_len)
{
	return 0 == CMBSM4EncryptWithECB(key, input, length, output, output_len);
}

bool DecWithSM4_ECB(BYTE key[16], BYTE* input, int length, BYTE* output, int* output_len)
{
	return 0 == CMBSM4DecryptWithECB(key, input, length, output, output_len);
}

bool SM3File(char* file, BYTE hash[32])
{
	if (0 == CMBSM3FileDigest(file, hash))
		return true;
	else
		return false;
}

RSA *GetServerRSAPubKey()
{
	const char PublicKeyData[] = "-----BEGIN PUBLIC KEY-----\n"
		"MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC9dYbXjQM77pdpddBS7bdEMVjT\n"
		"rMfvDyy3v+XEMAXBeh1j4M8670VoblikavFtX/xf9kP+byp3MtnnMqAhP0TkuUWw\n"
		"s8S157RuLaM8571TfHztPet4O23+2wTazpPuU/9ZVHDRxROMOCk2O6UYbEPMyQXv\n"
		"Ue2HA8c1ZRuP4+Q8RQIDAQAB\n"
		"-----END PUBLIC KEY-----";

	BIO	*bio = NULL;

	if( (bio=BIO_new_mem_buf((void *)PublicKeyData, sizeof(PublicKeyData))) == NULL)
		return NULL;

	RSA *rsa = PEM_read_bio_RSA_PUBKEY(bio, NULL, NULL, NULL);	
	BIO_free(bio);

	return rsa;
}

// 使用服务器公钥解密
bool DecWithServerPubKey(BYTE *pSource, int nSourceLen, BYTE *pDest, int *pDestLen)
{
	RSA *pkey = GetServerRSAPubKey();
	if (pkey == NULL)
		return false;


	return DecWithRsaPubKey(pSource, nSourceLen, pDest, pDestLen, pkey);
}