#include "stdio.h"
#include "BSCrypto.h"
#include "sha256.h"
#include "aria.h"

int _nState = BS_STATE_NONE;
BYTE _pEncKey[32] = { 0, };
BYTE _pDecKey[32] = { 0, };
ARIA_Key _AEncKey;
ARIA_Key _ADecKey;
SHA256_CTX _EncCTX;
SHA256_CTX _DecCTX;

// 문자열키를 32byte 바이너리 키로 생성, 모듈 내부에서만 사용됨
void _BuildStrToKey(const char* sKey, BYTE* pOut, int nSize)
{
	for (int i = 0; i < nSize; ++i)
		pOut[i] = i;

	int n = strlen(sKey);
	for (int i = 0; i < n; ++i)
	{
		pOut[i % nSize] ^= sKey[i];
	}
}

char* BSCrypto_GetVersion()
{
	return VER_BSCRYPTO;
}

// 암호모듈을 초기화하고 자가시험을 수행한다.
// BS_SUCCESS (0x0001) : 초기화 성공
// BS_CRYPTO_ERROR_SELFTEST_FAILED (0x0002) : 자가시험 실패
// BS_CRYPTO_ERROR_STATE_IN_ERROR (0x0003) : 오류 상태
int BSCrypto_Initialize()
{
	if (_nState != BS_STATE_NONE)
		return BS_SUCCESS;

// 자가시험
	if (BSCrypto_SelfTest() != BS_SUCCESS)
		return BS_ERROR_INIT;

// 초기화
	BSCrypto_Clear();
	_nState = BS_STATE_INIT;

	return BS_SUCCESS;
}

int BSCrypto_SelfTest()
{
	BYTE rk[16 * 17], c[16], * b, mk[32];
	int i, flag;
	const DWORD NUMBER = 0x00000042;
	BYTE p[16] = { 0x11, 0x11, 0x11, 0x11, 0xaa, 0xaa, 0xaa, 0xaa,
				  0x11, 0x11, 0x11, 0x11, 0xbb, 0xbb, 0xbb, 0xbb };
	const BYTE cryptResult[] = {
				  0x8d, 0x14, 0x70, 0x62, 0x5f, 0x59, 0xeb, 0xac,
				  0xb0, 0xe5, 0x5b, 0x53, 0x4b, 0x3e, 0x46, 0x2b };

// 리틀 엔디안 체크
	b = (BYTE*)(&NUMBER);
	if (b[0] != 0x42) 
		return BS_ERROR_CRYPTO_SELFTEST_FAILED; // BIG ENDIAN platform

	for (i = 0; i < 16; i++)
		mk[i] = i * 0x11;

	for (i = 16; i < 24; i++)
		mk[i] = (i - 16) * 0x11;

	Crypt(p, EncKeySetup(mk, rk, 192), rk, c);
	flag = 0;

	for (i = 0; i < 16; i++)
		if (c[i] != cryptResult[i])
			flag = 1;

	if (flag == 1)
		return BS_ERROR_CRYPTO_SELFTEST_FAILED;

	for (i = 0; i < 32; i++)
		mk[i] = 0;

	for (i = 0; i < 16; i++)
		p[i] = 0;

	EncKeySetup(mk, rk, 192);
	Crypt(p, 14, rk, c);
	DecKeySetup(mk, rk, 192);
	Crypt(c, 14, rk, p);
	flag = 0;
	for (i = 0; i < 16; i++)
		if (p[i] != 0)
			flag = 1;

	if (flag == 1)
		return BS_ERROR_CRYPTO_SELFTEST_FAILED;

	return BS_SUCCESS;
}

int BSCrypto_Clear()
{
	if (_nState != BS_STATE_INIT)
		return BS_ERROR_STATE_IN_ERROR;

	memset(&_pEncKey, 0, 32);
	memset(&_pDecKey, 0, 32);
	memset(&_AEncKey, 0, sizeof(_AEncKey));
	memset(&_ADecKey, 0, sizeof(_AEncKey));
	memset(&_EncCTX, 0, sizeof(_EncCTX));
	memset(&_DecCTX, 0, sizeof(_DecCTX));

	return BS_SUCCESS;
}

int BSCrypto_GetState()
{
	return _nState;
}

int BSCrypto_FileToSha256(LPCWCH sPath, BYTE* pOutHash)
{
	FILE* f = NULL;
	_wfopen_s(&f, sPath, L"rb");
	if (f != NULL)
	{
		fseek(f, 0, SEEK_END);
		LONGLONG llRead = 0, llReaded = 0;
		LONGLONG llSize = _ftelli64(f);
		fseek(f, 0, SEEK_SET);

		BYTE pBuf[BUF_SIZE] = { 0, };

		SHA256_CTX ctx;
		sha256_init(&ctx);
		while (llReaded < llSize)
		{
			if (llSize < (llReaded + BUF_SIZE))
				llRead = llSize - llReaded;
			else
				llRead = BUF_SIZE;

			llRead = fread_s(pBuf, BUF_SIZE, 1, llRead, f);
			if (llRead == 0)
				break;

			llReaded += llRead;

			sha256_update(&ctx, pBuf, llRead);
		}
		sha256_final(&ctx, pOutHash);
		fclose(f);
		return BS_SUCCESS;
	}
	else
		return BS_ERROR_OPEN_FILE;
}

void BSCrypto_SetRandomStr(char* sRnd, int nLen)
{
	int i;
	srand((unsigned int)time(NULL));
	for (i = 0; i < nLen; i++)
		sRnd[i] = 'a' + rand() % 26;
}

int BSCrypto_EncInit(const char* sKey)
{
	if (_nState != BS_STATE_INIT)
		return BS_ERROR_STATE_IN_ERROR;

	sha256_init(&_EncCTX);
	_BuildStrToKey(sKey, _pEncKey, 32);
	ARIA_EncKeySetup(_pEncKey, &_AEncKey, 256);
	_nState = BS_STATE_ENC;

	return BS_SUCCESS;
}

int BSCrypto_EncUpdate(const BYTE* pIn, BYTE* pOut, int nIn)
{
	if (_nState != BS_STATE_ENC)
		return BS_ERROR_STATE_IN_ERROR;

	sha256_update(&_EncCTX, pIn, nIn);
	ARIA_Crypt(pIn, pOut, nIn, &_AEncKey);

	return BS_SUCCESS;
}

int BSCrypto_EncFinal(BYTE** pOutHash)
{
	if (_nState != BS_STATE_ENC)
		return BS_ERROR_STATE_IN_ERROR;

	BYTE pHash[SHA256_BLOCK_SIZE];
	sha256_final(&_EncCTX, pHash);
	if (pOutHash != NULL)
		memcpy(*pOutHash, pHash, SHA256_BLOCK_SIZE);

	_nState = BS_STATE_INIT;
	BSCrypto_Clear;

	return BS_SUCCESS;
}

int BSCrypto_DecInit(const char* sKey)
{
	if (_nState != BS_STATE_INIT)
		return BS_ERROR_STATE_IN_ERROR;

	sha256_init(&_DecCTX);
	_BuildStrToKey(sKey, _pDecKey, 32);
	ARIA_EncKeySetup(_pDecKey, &_ADecKey, 256);
	_nState = BS_STATE_DEC;

	return BS_SUCCESS;
}

int BSCrypto_DecUpdate(const BYTE* pIn, BYTE* pOut, int nIn)
{
	if (_nState != BS_STATE_DEC)
		return BS_ERROR_STATE_IN_ERROR;

	sha256_update(&_DecCTX, pIn, nIn);
	ARIA_Crypt(pIn, pOut, nIn, &_ADecKey);

	return BS_SUCCESS;
}

int BSCrypto_DecFinal(BYTE** pOutHash)
{
	if (_nState != BS_STATE_DEC)
		return BS_ERROR_STATE_IN_ERROR;

	BYTE pHash[SHA256_BLOCK_SIZE];
	sha256_final(&_DecCTX, pHash);
	if (pOutHash != NULL)
		memcpy(*pOutHash, pHash, SHA256_BLOCK_SIZE);

	_nState = BS_STATE_INIT;
	BSCrypto_Clear;

	return BS_SUCCESS;
}

int BSCrypto_EncFile(const char* sKey, LPCWCH sSrcPath, LPCWCH sDestPath)
{
	BSCryptoHead HeadInfo;
	memset(&HeadInfo, 0, sizeof(HeadInfo));
	HeadInfo.nHeadLen = sizeof(HeadInfo);

	int nPad = 0;
	FILE* fSrc = NULL;
	_wfopen_s(&fSrc, sSrcPath, L"rb");
	if (fSrc != NULL)
	{
		fseek(fSrc, 0, SEEK_END);
		LONGLONG llRead = 0, llReaded = 0;
		LONGLONG llSize = _ftelli64(fSrc);
		fseek(fSrc, 0, SEEK_SET);

		BYTE pBuf[BUF_SIZE] = { 0, };
		BYTE pBufOut[BUF_SIZE] = { 0, };
		BYTE pHash[SHA256_BLOCK_SIZE];

		FILE* fDest = NULL;
		_wfopen_s(&fDest, sDestPath, L"wb");
		if (fDest == NULL)
			return BS_ERROR_OPEN_FILE;

	// 시그니처 추가
		fwrite(ENC_SIG, 1, 5, fDest);

	// 헤더정보 채울 수 있는거 채우고
		HeadInfo.llSize = llSize;
		time(&HeadInfo.tTime);
	 	BSCrypto_SetRandomStr(HeadInfo.sReserve, 220);

	// 헤더자리 비우고 시작
		fseek(fDest, HeadInfo.nHeadLen, SEEK_CUR);

	// 암호화 준비
		BYTE pKey[32];
		_BuildStrToKey(sKey, pKey, 32);
		ARIA_Key AKey;
		ARIA_EncKeySetup(pKey, &AKey, 256);

		SHA256_CTX ctx;
		sha256_init(&ctx);
		while (llReaded < llSize)
		{
			if (llSize < (llReaded + BUF_SIZE))
			{
				llRead = llSize - llReaded;
				nPad = llRead % 16;
				if (nPad > 0)
				{
					nPad = 16 - nPad;
					memset(pBufOut, 0, BUF_SIZE);
				}
			}
			else
			{
				llRead = BUF_SIZE;
			}

			llRead = fread_s(pBuf, BUF_SIZE, 1, llRead, fSrc);
			if (llRead == 0)
				break;

			llReaded += llRead;

			ARIA_Crypt(pBuf, pBufOut, llRead + nPad, &AKey);
			fwrite(pBufOut, 1, llRead + nPad, fDest);

			sha256_update(&ctx, pBuf, llRead);
		}
		sha256_final(&ctx, pHash);
	// 원본 해시값 입력
		memcpy(HeadInfo.pOrgHash, pHash, SHA256_BLOCK_SIZE);

	// 헤더 정보 암호화 --------------------------------------------------		
		_BuildStrToKey(PASS_HEAD, pKey, 32);
		ARIA_EncKeySetup(pKey, &AKey, 256);
		ARIA_Crypt((BYTE*)&HeadInfo, pBufOut, HeadInfo.nHeadLen, &AKey);		
	// 헤더 정보 암호화 --------------------------------------------------

	// 암호화된 헤더 쓰기
		fseek(fDest, 5, SEEK_SET);
		fwrite(pBufOut, 1, HeadInfo.nHeadLen, fDest);

		if (fDest != NULL)
			fclose(fDest);
		fclose(fSrc);

		return BS_SUCCESS;
	} else 
		return BS_ERROR_OPEN_FILE;
}

int BSCrypto_DecFile(const char* sKey, LPCWCH sSrcPath, LPCWCH sDestPath)
{
	FILE* fEnc = NULL;
	_wfopen_s(&fEnc, sSrcPath, L"rb");
	if (fEnc != NULL)
	{
		fseek(fEnc, 0, SEEK_END);
		LONGLONG llSize = _ftelli64(fEnc);
		fseek(fEnc, 0, SEEK_SET);

		char sSig[5];
		if (fread_s(sSig, 5, 1, 5, fEnc) != 5)
			return BS_ERROR_INVALID_ENC_FILE;

		if (memcmp(sSig, ENC_SIG, 5) != 0)
			return BS_ERROR_INVALID_ENC_FILE;

		BYTE pBuf[BUF_SIZE] = { 0, };
		BYTE pBufOut[BUF_SIZE] = { 0, };
		BYTE pHash[SHA256_BLOCK_SIZE];

		BSCryptoHead HeadInfo;
		int nHeadLen = sizeof(HeadInfo);		
		memset(&HeadInfo, 0, nHeadLen);

	// 헤더 정보 복호화 --------------------------------------------------
		if (fread_s(pBuf, BUF_SIZE, 1, nHeadLen, fEnc) != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;

		BYTE pKey[32];
		_BuildStrToKey(PASS_HEAD, pKey, 32);

		ARIA_Key AKey;
		ARIA_DecKeySetup(pKey, &AKey, 256);
		ARIA_Crypt(pBuf, (BYTE*)&HeadInfo, nHeadLen, &AKey);

		if (HeadInfo.nHeadLen != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;
	// 헤더 정보 복호화 --------------------------------------------------
		
	// 복호화 준비
		_BuildStrToKey(sKey, pKey, 32);
		ARIA_DecKeySetup(pKey, &AKey, 256);

		FILE* fDec = NULL;
		_wfopen_s(&fDec, sDestPath, L"wb");
		if (fDec == NULL)
			return BS_ERROR_CREATE_FILE;

		LONGLONG llWrite = 0, llRead = 0, llReaded = 0;
		SHA256_CTX ctx;
		sha256_init(&ctx);
		while (llReaded < llSize)
		{
			if (llSize < (llReaded + BUF_SIZE))
			{
				llRead = llSize - llReaded;
			// 원본 크기로 최종 보정
				llWrite = HeadInfo.llSize - llReaded;
			}
			else
			{
				llRead = BUF_SIZE;
				llWrite = llRead;
			}

			llRead = fread_s(pBuf, BUF_SIZE, 1, llRead, fEnc);
			if (llRead == 0)
				break;

			llReaded += llRead;

			ARIA_Crypt(pBuf, pBufOut, llRead, &AKey);
			fwrite(pBufOut, 1, llWrite, fDec);

			sha256_update(&ctx, pBufOut, llWrite);
		}
		sha256_final(&ctx, pHash);

		if (fDec != NULL)
			fclose(fDec);
		fclose(fEnc);

	// 원본과 해시 비교 후 종료
		if (memcmp(HeadInfo.pOrgHash, pHash, SHA256_BLOCK_SIZE) != 0)
			return BS_ERROR_DECRYPT_FILE;
		else
			return BS_SUCCESS;
	} else
		return BS_ERROR_OPEN_FILE;
}

int BSCrypto_GetOrgFileHash(LPCWCH sEncPath, BYTE* pOutHash)
{
	FILE* fEnc = NULL;
	_wfopen_s(&fEnc, sEncPath, L"rb");
	if (fEnc != NULL)
	{
		fseek(fEnc, 0, SEEK_SET);

		char sSig[5];
		if (fread_s(sSig, 5, 1, 5, fEnc) != 5)
			return BS_ERROR_INVALID_ENC_FILE;

		if (memcmp(sSig, ENC_SIG, 5) != 0)
			return BS_ERROR_INVALID_ENC_FILE;

		BYTE pBuf[BUF_SIZE] = { 0, };

		BSCryptoHead HeadInfo;
		int nHeadLen = sizeof(HeadInfo);
		memset(&HeadInfo, 0, nHeadLen);

		// 헤더 정보 복호화 --------------------------------------------------
		if (fread_s(pBuf, BUF_SIZE, 1, nHeadLen, fEnc) != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;
		fclose(fEnc);

		BYTE pKey[32];
		_BuildStrToKey(PASS_HEAD, pKey, 32);

		ARIA_Key AKey;
		ARIA_DecKeySetup(pKey, &AKey, 256);
		ARIA_Crypt(pBuf, (BYTE*)&HeadInfo, nHeadLen, &AKey);

		if (HeadInfo.nHeadLen != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;
		// 헤더 정보 복호화 --------------------------------------------------

		memcpy(pOutHash, HeadInfo.pOrgHash, SHA256_BLOCK_SIZE);

		return BS_SUCCESS;
	}
	else
		return BS_ERROR_OPEN_FILE;
}

int BSCrypto_IsEncFile(LPCWCH sEncPath)
{
	FILE* fEnc = NULL;
	_wfopen_s(&fEnc, sEncPath, L"rb");
	if (fEnc != NULL)
	{
		fseek(fEnc, 0, SEEK_SET);

		char sSig[5];
		if (fread_s(sSig, 5, 1, 5, fEnc) != 5)
			return BS_ERROR_INVALID_ENC_FILE;

		if (memcmp(sSig, ENC_SIG, 5) != 0)
			return BS_ERROR_INVALID_ENC_FILE;

		BYTE pBuf[BUF_SIZE] = { 0, };

		BSCryptoHead HeadInfo;
		int nHeadLen = sizeof(HeadInfo);
		memset(&HeadInfo, 0, nHeadLen);

		// 헤더 정보 복호화 --------------------------------------------------
		if (fread_s(pBuf, BUF_SIZE, 1, nHeadLen, fEnc) != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;
		fclose(fEnc);

		BYTE pKey[32];
		_BuildStrToKey(PASS_HEAD, pKey, 32);

		ARIA_Key AKey;
		ARIA_DecKeySetup(pKey, &AKey, 256);
		ARIA_Crypt(pBuf, (BYTE*)&HeadInfo, nHeadLen, &AKey);

		if (HeadInfo.nHeadLen != nHeadLen)
			return BS_ERROR_INVALID_ENC_FILE;
		// 헤더 정보 복호화 --------------------------------------------------

		return BS_SUCCESS;
	}
	else
		return BS_ERROR_OPEN_FILE;
}