#include "stdafx.h"
#include "modVer.h"
#include "SpBase.h"

#include "TokenKeeper_server_g.h"

using namespace TokenKeeper;

#include "RVCComm.h"

class CTokenManagerEntity;

class CTokenManagerSession : public TokenService_ServerSessionBase
{
public:
	CTokenManagerSession(CTokenManagerEntity *pEntity) : m_pEntity(pEntity) { }
	virtual ~CTokenManagerSession() { }

	virtual void Handle_GetToken(SpReqAnsContext<TokenService_GetToken_Req, TokenService_GetToken_Ans>::Pointer ctx);

	virtual void Handle_SetToken(SpReqAnsContext<TokenService_SetToken_Req, TokenService_SetToken_Ans>::Pointer ctx);

	virtual void Handle_GetSharedSK(SpReqAnsContext<TokenService_GetSharedSK_Req, TokenService_GetSharedSK_Ans>::Pointer ctx);

	virtual void Handle_SetSharedSK(SpReqAnsContext<TokenService_SetSharedSK_Req, TokenService_SetSharedSK_Ans>::Pointer ctx);

	virtual void Handle_SetKeyPair(SpReqAnsContext<TokenService_SetKeyPair_Req, TokenService_SetKeyPair_Ans>::Pointer ctx);

	virtual void Handle_GetPublicKey(SpReqAnsContext<TokenService_GetPublicKey_Req, TokenService_GetPublicKey_Ans>::Pointer ctx);

	virtual void Handle_EncryptWithPubKey(SpReqAnsContext<TokenService_EncryptWithPubKey_Req, TokenService_EncryptWithPubKey_Ans>::Pointer ctx);

	virtual void Handle_DecryptWithPubKey(SpReqAnsContext<TokenService_DecryptWithPubKey_Req, TokenService_DecryptWithPubKey_Ans>::Pointer ctx);

	virtual void Handle_EncryptWithPriKey(SpReqAnsContext<TokenService_EncryptWithPriKey_Req, TokenService_EncryptWithPriKey_Ans>::Pointer ctx);

	virtual void Handle_DecryptWithPriKey(SpReqAnsContext<TokenService_DecryptWithPriKey_Req, TokenService_DecryptWithPriKey_Ans>::Pointer ctx);

private:
	CTokenManagerEntity *m_pEntity;
};


class CTokenManagerEntity : public CEntityBase
{
public:

	CTokenManagerEntity() : m_bTokenInitialized(false), m_bKeyInitialized(false), m_bSSKInitialized(false) {}
	virtual ~CTokenManagerEntity() {}

	virtual const char *GetEntityName() const { return "TokenKeeper"; }
	const char* GetEntityVersion() const { return MODULE_VERSION_FULL; }

	virtual bool IsService()const { return true; }

	virtual void OnStarted()
	{
		InitializeCriticalSection(&m_lock);
	}

	virtual void OnPreClose(EntityCloseCauseEnum eCloseCause,CSmartPointer<ITransactionContext> pTransactionContext)
	{
		DeleteCriticalSection(&m_lock);
		pTransactionContext->SendAnswer(Error_Succeed);
	}

	virtual CServerSessionBase *OnNewSession(const char* pszRemoteEntityName, const char * pszClass)
	{
		Dbg("new session required from %s", pszRemoteEntityName);
		return new CTokenManagerSession(this);
	}

	ErrorCodeEnum GetToken(CBlob &token)
	{
		ErrorCodeEnum Error;
		EnterCriticalSection(&m_lock);
		if (m_bTokenInitialized) {
			token = m_blobToken;
			Error = Error_Succeed;
		} else {
			Error = Error_NotInit;
			Dbg("%s: token is not initiazlied!", __FUNCTION__);
		}
		LeaveCriticalSection(&m_lock);
		return Error;
	}

	ErrorCodeEnum SetToken(const CBlob &token)
	{
		EnterCriticalSection(&m_lock);
		m_blobToken = token;
		m_bTokenInitialized = true;
		LeaveCriticalSection(&m_lock);
		return Error_Succeed;
	}

	ErrorCodeEnum GetSharedSK(CBlob &ssk)
	{
		ErrorCodeEnum Error;
		EnterCriticalSection(&m_lock);
		if (m_bSSKInitialized) {
			ssk = m_blobSSK;
			Error = Error_Succeed;
		} else {
			Error = Error_NotInit;
            Dbg("%s: SharedSK is not initiazlied!", __FUNCTION__);
		}
		LeaveCriticalSection(&m_lock);
		return Error;
	}

	ErrorCodeEnum SetSharedSK(const CBlob &ssk)
	{
		EnterCriticalSection(&m_lock);
		m_blobSSK = ssk;
		m_bSSKInitialized = true;
		LeaveCriticalSection(&m_lock);
		return Error_Succeed;
	}

	ErrorCodeEnum SetKeyPair(const CBlob &pri, const CBlob &pub)
	{
		EnterCriticalSection(&m_lock);
		m_blobPriKey = pri;
		m_blobPubKey = pub;
		m_bKeyInitialized = true;
		LeaveCriticalSection(&m_lock);
		return Error_Succeed;
	}

	ErrorCodeEnum GetPubKey(CBlob &pub)
	{
		ErrorCodeEnum Error;
		EnterCriticalSection(&m_lock);
		if (m_bKeyInitialized) {
			pub = m_blobPubKey;
			Error = Error_Succeed;
		} else {
			LogError(Severity_Low, Error_NotInit, 0, "key not initialized");
			Error = Error_NotInit;
		}
		LeaveCriticalSection(&m_lock);
		return Error;
	}

	ErrorCodeEnum EncryptWithPubKey(const CBlob &raw_data, CBlob &enc_data)
	{
		ErrorCodeEnum Error;
		EnterCriticalSection(&m_lock);
		if (!m_bKeyInitialized) {
			LogError(Severity_Low, Error_NotInit, 0, "key not initialized");
			Error = Error_NotInit;
		}
		else
		{
			// 保证目标缓冲区长度
			int len = raw_data.m_iLength;
			if (len % 117 == 0)
				len = (len / 117) * 128;
			else
				len = (len / 117 + 1) * 128;

			enc_data.Alloc(len + 4);
			bool bRet = EncWithSM2PubKey((BYTE*)raw_data.m_pData, raw_data.m_iLength, (BYTE*)enc_data.m_pData + 4, &len, (BYTE*)m_blobPubKey.m_pData, m_blobPubKey.m_iLength);
			enc_data.Resize(len + 4);		// shrink to return size
			auto dstData = (byte*)enc_data.m_pData;
			dstData[3] = (byte)((len >> 24) & 0xFF);
			dstData[2] = (byte)((len >> 16) & 0xFF);
			dstData[1] = (byte)((len >> 8) & 0xFF);
			dstData[0] = (byte)(len & 0xFF);


			Error = bRet ? Error_Succeed : Error_Unexpect;
		}
		LeaveCriticalSection(&m_lock);
		return Error;
	}

	ErrorCodeEnum DecryptWithPubKey(const CBlob &enc_data, CBlob &raw_data)
	{
		/*
		ErrorCodeEnum Error(Error_Succeed);
		EnterCriticalSection(&m_lock);
		if (!m_bKeyInitialized) {
			LogError(Severity_Low, Error_NotInit, 0, "key not initialized");
			Error = Error_NotInit;
		} else {
			int len = enc_data.m_iLength;
			raw_data.Alloc(len);
			bool bRet = DecWithRsaPubKey((BYTE*)enc_data.m_pData, enc_data.m_iLength, (BYTE*)raw_data.m_pData, &len, (BYTE*)m_blobPubKey.m_pData, m_blobPubKey.m_iLength);
			raw_data.Resize(len);		// shrink to return size
            if (!bRet) {
                Error = Error_DataCheck;
                LogWarn(Severity_Low, Error, 0, "DecWithRsaPubKey failed!");
            }
		}
		LeaveCriticalSection(&m_lock);
		*/
		LogWarn(Severity_Low, Error_Cancel, 0, "DecWithRsaPubKey not support in sm os!");

		return Error_Cancel;
	}

	ErrorCodeEnum EncryptWithPriKey(const CBlob &raw_data, CBlob &enc_data)
	{
		/*
		ErrorCodeEnum Error(Error_Succeed);
		EnterCriticalSection(&m_lock);
		if (!m_bKeyInitialized) {
			LogError(Severity_Low, Error_NotInit, 0, "key not initialized");
			Error = Error_NotInit; 
		} 
		else
		{
			// 保证目标缓冲区长度
			int len = raw_data.m_iLength;
			if (len % 117 ==0)
				len = (len / 117) * 128;
			else
				len = (len / 117 + 1) * 128;

			enc_data.Alloc(len);
			bool bRet = EncWithRsaPriKey((BYTE*)raw_data.m_pData, raw_data.m_iLength, (BYTE*)enc_data.m_pData, &len, (BYTE*)m_blobPriKey.m_pData, m_blobPriKey.m_iLength);
			enc_data.Resize(len);		// shrink to return size
            if (!bRet) {
                Error = Error_DataCheck;
                LogWarn(Severity_Low, Error, 0, "EncWithRsaPriKey failed!");
            }
		}
		LeaveCriticalSection(&m_lock);
		*/
		LogWarn(Severity_Low, Error_Cancel, 0, "EncryptWithPriKey not support in sm os!");
		return Error_Cancel;
	}
	ErrorCodeEnum DecryptWithPriKey(const CBlob &enc_data, CBlob &raw_data)
	{
		ErrorCodeEnum Error;
		EnterCriticalSection(&m_lock);
		if (!m_bKeyInitialized)
		{
			LogError(Severity_Low, Error_NotInit, 0, "key not initialized");
			Error = Error_NotInit;
		}
		else if (enc_data.m_iLength < 4)
		{
			LogError(Severity_Low, Error_NotInit, 0, "DecryptWithPriKey len < 4");
			Error = Error_NotInit;
		}
		else
		{
			auto srcData = (byte*)enc_data.m_pData;
			int len = ((int)srcData[3]) << 24 | ((int)srcData[2]) << 16 | ((int)srcData[1]) << 8 | ((int)srcData[0]);
			raw_data.Alloc(len);
			bool bRet = DecWithSM2PriKey((BYTE*)enc_data.m_pData + 4, len, (BYTE*)raw_data.m_pData, &len, (BYTE*)m_blobPriKey.m_pData, m_blobPriKey.m_iLength);
			raw_data.Resize(len);		// shrink to return size
			Error = bRet ? Error_Succeed : Error_Unexpect;
		}
		LeaveCriticalSection(&m_lock);
		return Error;
	}

private:
	CBlob m_blobToken;
	bool m_bTokenInitialized;
	CBlob m_blobSSK;
	bool m_bSSKInitialized;
	CBlob m_blobPriKey;
	CBlob m_blobPubKey;
	bool m_bKeyInitialized;
	CRITICAL_SECTION m_lock;
};

void CTokenManagerSession::Handle_GetToken( SpReqAnsContext<TokenService_GetToken_Req, TokenService_GetToken_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->GetToken(ctx->Ans.token);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_SetToken( SpReqAnsContext<TokenService_SetToken_Req, TokenService_SetToken_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->SetToken(ctx->Req.token);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_GetSharedSK(SpReqAnsContext<TokenService_GetSharedSK_Req, TokenService_GetSharedSK_Ans>::Pointer ctx)
{
	ErrorCodeEnum Error = m_pEntity->GetSharedSK(ctx->Ans.ssk);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_SetSharedSK(SpReqAnsContext<TokenService_SetSharedSK_Req, TokenService_SetSharedSK_Ans>::Pointer ctx)
{
	ErrorCodeEnum Error = m_pEntity->SetSharedSK(ctx->Req.ssk);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_SetKeyPair( SpReqAnsContext<TokenService_SetKeyPair_Req, TokenService_SetKeyPair_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->SetKeyPair(ctx->Req.pri_key, ctx->Req.pub_key);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_GetPublicKey( SpReqAnsContext<TokenService_GetPublicKey_Req, TokenService_GetPublicKey_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->GetPubKey(ctx->Ans.key);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_EncryptWithPubKey( SpReqAnsContext<TokenService_EncryptWithPubKey_Req, TokenService_EncryptWithPubKey_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->EncryptWithPubKey(ctx->Req.raw_data, ctx->Ans.enc_data);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_DecryptWithPubKey( SpReqAnsContext<TokenService_DecryptWithPubKey_Req, TokenService_DecryptWithPubKey_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->DecryptWithPubKey(ctx->Req.enc_data, ctx->Ans.raw_data);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_EncryptWithPriKey( SpReqAnsContext<TokenService_EncryptWithPriKey_Req, TokenService_EncryptWithPriKey_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->EncryptWithPriKey(ctx->Req.raw_data, ctx->Ans.enc_data);
	ctx->Answer(Error);
}

void CTokenManagerSession::Handle_DecryptWithPriKey( SpReqAnsContext<TokenService_DecryptWithPriKey_Req, TokenService_DecryptWithPriKey_Ans>::Pointer ctx )
{
	ErrorCodeEnum Error = m_pEntity->DecryptWithPriKey(ctx->Req.enc_data, ctx->Ans.raw_data);
	ctx->Answer(Error);
}

SP_BEGIN_ENTITY_MAP()
	SP_ENTITY(CTokenManagerEntity)
SP_END_ENTITY_MAP()