#include "stdafx.h"
#include "filecryption.h"
#include "fileanalysis.h"
#include "rvcfileheader.h"
#include "asf.h"

#include <iostream>
using namespace std;

#include <stdio.h>

#ifndef RVC_MAX_HEADER_LEN
#define RVC_MAX_HEADER_LEN 512
#endif

#ifndef RVC_FILE_INT_LEN
#define RVC_FILE_INT_LEN 4
#endif

#ifndef RVC_READ_BUFFER_SIZE
#define RVC_READ_BUFFER_SIZE 1024
#endif


#ifndef safe_free
#define safe_free(x) if (x){ free(x); x = NULL;}
#endif

int encryption_file(char* poutfile, size_t uoutlen, const char* pfilename, const filecryption_callback_t* pcallback, eRvcCryptionVersion eversion)
{
	int iret = -1;

	if (NULL == pfilename){
		return iret;
	}

	char pbuffer[RVC_MAX_HEADER_LEN] = {0};
	size_t uheaderlen = RVC_MAX_HEADER_LEN;
	int ilen = constrcut_rvc_file_header(pbuffer, &uheaderlen, pfilename, pcallback, eversion);
	if (0 != ilen){
		return iret;
	}

	unsigned char strkey[SM4ENC_BLOCK_SIZE] = {0};
	size_t uindex = (RVC_FILE_HEADER_FLAG_LEN + sizeof(size_t));
	if (get_key_from_header_info(strkey, SM4ENC_BLOCK_SIZE, pbuffer + uindex, uheaderlen - uindex, pcallback)){
		return iret;
	}

	if (get_encrytion_filename(poutfile, uoutlen, false, RVC_FILEENC_STR, strlen(RVC_FILEENC_STR), pfilename)){
		return iret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		return iret;
	}

	FILE* pDestFile = fopen(poutfile, "wb");
	if (NULL == pDestFile){
		fclose(pSrcFile);
		return iret;
	}
	
	//1. 新增自定义文件头
	fwrite(pbuffer, 1, uheaderlen, pDestFile);

	//获取加密前文件头GUID
	char strGuid[HEADER_OBJECT_GUID_LEN] = {0};
	fread(strGuid, 1, HEADER_OBJECT_GUID_LEN, pSrcFile);
	if (0 != memcmp(strGuid, ASF_Header_GUID, HEADER_OBJECT_GUID_LEN)){
		safe_log(pcallback, "%s", "math guid failed, current format is not surpport!");
		fclose(pSrcFile);
		fclose(pDestFile);
		return iret;
	}

	//2. 拷贝原文件头guid
	fwrite(strGuid, 1, HEADER_OBJECT_GUID_LEN, pDestFile);

	char strheadsize[ASF_HEADER_SIZE_LEN] = {0};
	fread(strheadsize, 1, ASF_HEADER_SIZE_LEN, pSrcFile);
	uint64_t uheadlen = get_asf_headsize(strheadsize, ASF_HEADER_SIZE_LEN);

	uint64_t uheadobjlen = uheadlen - HEADER_OBJECT_GUID_LEN - ASF_HEADER_SIZE_LEN;
	unsigned char* pobjdata = (unsigned char*)malloc(uheadobjlen);
	fread(pobjdata, 1, uheadobjlen, pSrcFile);

	int iencdatalen = uheadobjlen + SM4ENC_BLOCK_SIZE;
	unsigned char* pobjdataenc = (unsigned char*)malloc(iencdatalen);
	SM4EncECBMode(strkey, pobjdata, uheadobjlen, pobjdataenc, &iencdatalen);

	char strdatalen[ASF_HEADER_SIZE_LEN] = {0};
	construct_asf_headsize(strdatalen, ASF_HEADER_SIZE_LEN, iencdatalen + HEADER_OBJECT_GUID_LEN + ASF_HEADER_SIZE_LEN);
	//3. 填充加密后文件头长度
	fwrite(strdatalen, 1, ASF_HEADER_SIZE_LEN, pDestFile);

	//4. 填充加密后文件头内容
	fwrite(pobjdataenc, 1, iencdatalen, pDestFile);

	char *data_buffer=(char *)malloc(RVC_READ_BUFFER_SIZE);
	size_t udatalen = 0;
	while(udatalen = fread(data_buffer, 1, RVC_READ_BUFFER_SIZE, pSrcFile)){
		fwrite(data_buffer, 1, udatalen, pDestFile);	
	}

	fclose(pSrcFile);
	fclose(pDestFile);

	safe_free(pobjdata);
	safe_free(pobjdataenc);
	safe_free(data_buffer);

	iret = 0;

	return iret;
}


int decryption_file(char* poutfile, size_t uoutlen, const char* pfilename, const filecryption_callback_t* pcallback, eRvcCryptionVersion eversion)
{
	int iret = -1;

	if (NULL == pfilename || eversion >= sizeof(cryption_ver_flag_table)/sizeof(char*) || eversion < 0){
		safe_log(pcallback,"%s","invalid encryption version param.");
		return iret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		safe_log(pcallback,"open file %s failed!",pfilename);
		return iret;
	}

	char strrvcflag[RVC_FILE_HEADER_FLAG_LEN] = {0};
	fread(strrvcflag,1, RVC_FILE_HEADER_FLAG_LEN, pSrcFile);
	if (0 != memcmp(strrvcflag, rvc_header, RVC_FILE_HEADER_FLAG_LEN)){
		safe_log(pcallback,"file %s is not encryption!",pfilename);
		fclose(pSrcFile);
		iret = 0;
		return iret;
	}

	int irvcheadlen = 0;
	fread(&irvcheadlen, 1, sizeof(int), pSrcFile);
	char* prvcbuffer = (char*)malloc(irvcheadlen);
	memset(prvcbuffer, 0, irvcheadlen);
	fread(prvcbuffer, 1, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(int), pSrcFile);

	if (0 != memcmp(prvcbuffer+sizeof(size_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN)){
		if (0 == memcmp(prvcbuffer+sizeof(size_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN-1)){
			safe_log(pcallback,"file %s encryption and decryption version is not matched!",pfilename);
		} 
		fclose(pSrcFile);
		return iret;
	}
	
	unsigned char strkey[SM4ENC_BLOCK_SIZE] = {0};
	if (get_key_from_header_info(strkey, SM4ENC_BLOCK_SIZE, prvcbuffer, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(int), pcallback)){
		return iret;
	}

	if (get_decrytion_filename(poutfile, uoutlen, RVC_FILEDEC_STR, strlen(RVC_FILEDEC_STR), pfilename)){
		return iret;
	}

	FILE* pDestFile = fopen(poutfile, "wb");
	if (NULL == pDestFile){
		return iret;
	}

	char strasfheader[ASF_HEADER_GUID_LEN] = {0};
	if (ASF_HEADER_GUID_LEN == fread(strasfheader, 1, ASF_HEADER_GUID_LEN, pSrcFile)){
		fwrite(strasfheader, 1, ASF_HEADER_GUID_LEN, pDestFile);
	}

	char strheadlen[ASF_HEADER_SIZE_LEN] = {0};
	fread(strheadlen, 1, ASF_HEADER_SIZE_LEN, pSrcFile);
	uint64_t uheadlen = get_asf_headsize(strheadlen, ASF_HEADER_SIZE_LEN);
	unsigned char* pheaddata = (unsigned char*)malloc(uheadlen);
	unsigned char* pdecheaddata = (unsigned char*)malloc(uheadlen);
	int ioutlen = uheadlen;

	fread(pheaddata, 1, uheadlen - ASF_HEADER_GUID_LEN - ASF_HEADER_SIZE_LEN, pSrcFile);
	SM4DecECBMode(strkey, pheaddata, uheadlen - ASF_HEADER_GUID_LEN - ASF_HEADER_SIZE_LEN, pdecheaddata, &ioutlen);
	construct_asf_headsize(strheadlen,ASF_HEADER_SIZE_LEN,ioutlen + ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN);
	fwrite(strheadlen, 1, ASF_HEADER_SIZE_LEN, pDestFile);
	fwrite(pdecheaddata, 1, ioutlen, pDestFile);

	char *data_buffer=(char *)malloc(RVC_READ_BUFFER_SIZE);
	size_t udatalen = 0;
	while(udatalen = fread(data_buffer, 1, RVC_READ_BUFFER_SIZE, pSrcFile)){
		fwrite(data_buffer, 1, udatalen, pDestFile);	
	}

	fclose(pSrcFile);
	fclose(pDestFile);

	safe_free(data_buffer);
	safe_free(pdecheaddata);
	safe_free(pheaddata);

	return iret;
}


bool is_file_encrypted(const char* pfilename, const filecryption_callback_t* pcallback)
{
	bool bret = false;
	if (NULL == pfilename){
		return bret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		safe_log(pcallback,"open file %s failed!",pfilename);
		return bret;
	}

	char strrvcheader[RVC_FILE_HEADER_FLAG_LEN] = {0};
	if (RVC_FILE_HEADER_FLAG_LEN == fread(strrvcheader, 1, RVC_FILE_HEADER_FLAG_LEN, pSrcFile)){
		if (0 == strnicmp(strrvcheader, rvc_header, RVC_FILE_HEADER_FLAG_LEN)){
			bret = true;
		} 
		else{
			safe_log(pcallback,"file %s is not encryption!",pfilename);
		}
	}
	fclose(pSrcFile);

	return bret;
}


/*
文件解密功能；
	输入：文件路径
	输出：文件明文头数据，文件明文头长度，加解密前后文件偏移量，自定义头文件头信息
*/
int rvc_file_decrypt(unsigned char** pdechead, size_t* udecheadlen, int* ioffset, char** pstrjson, size_t* ujsonlen, const char* pfilename, const filecryption_callback_t* pcallback, eRvcCryptionVersion eversion)
{
	int iret = -1;
	if (NULL == pfilename || eversion >= sizeof(cryption_ver_flag_table)/sizeof(char*) || eversion < 0){
		safe_log(pcallback,"%s","invalid cryption version param.");
		return iret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		safe_log(pcallback,"open file %s failed!",pfilename);
		return iret;
	}

	unsigned char strkey[SM4ENC_BLOCK_SIZE] = {0};
	int irvcheadlen = 0;

	char strrvcflag[RVC_FILE_HEADER_FLAG_LEN] = {0};
	fread(strrvcflag,1, RVC_FILE_HEADER_FLAG_LEN, pSrcFile);
	if (0 != memcmp(strrvcflag, rvc_header, RVC_FILE_HEADER_FLAG_LEN)){
		//safe_log(pcallback,"file %s is not encrypted!",pfilename);
		*pdechead = NULL;
		*udecheadlen = 0;
		*ioffset = 0;
		fclose(pSrcFile);
		iret = 0;
		return iret;
	}
	else{
		//safe_log(pcallback,"file %s is encrypted!",pfilename);
		fread(&irvcheadlen, 1, sizeof(int), pSrcFile);
		char* prvcbuffer = (char*)malloc(irvcheadlen);
		memset(prvcbuffer, 0, irvcheadlen);
		fread(prvcbuffer, 1, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(int), pSrcFile);

		if (0 != memcmp(prvcbuffer+sizeof(size_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN)){
			if (0 == memcmp(prvcbuffer+sizeof(size_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN-1)){
				safe_log(pcallback,"file %s encryption and decrption version is not matched!",pfilename);
			} 
			safe_free(prvcbuffer);
			fclose(pSrcFile);
			return iret;
		}

		if (get_key_from_header_info(strkey, SM4ENC_BLOCK_SIZE, prvcbuffer, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(int), pcallback)){
			safe_free(prvcbuffer);
			fclose(pSrcFile);
			return iret;
		}

		int ijson = get_file_json_infos_from_rvc_header(pstrjson, ujsonlen, prvcbuffer, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(int));
		if (0 == ijson){
			safe_log(pcallback,"get_file_json_infos_from_rvc_header success!");
		}
		safe_free(prvcbuffer);
	}

	char strasfheader[ASF_HEADER_GUID_LEN] = {0};
	if (ASF_HEADER_GUID_LEN == fread(strasfheader, 1, ASF_HEADER_GUID_LEN, pSrcFile)){
		if (0 != memcmp(strasfheader, ASF_Header_GUID, ASF_HEADER_GUID_LEN)){
			safe_log(pcallback,"file %s is current not surrport format!",pfilename);
			fclose(pSrcFile);
			return iret;
		}
	}
	else{
		fclose(pSrcFile);
		return iret;
	}

	char strheadlen[ASF_HEADER_SIZE_LEN] = {0};
	fread(strheadlen, 1, ASF_HEADER_SIZE_LEN, pSrcFile);
	uint64_t uheadlen = get_asf_headsize(strheadlen, ASF_HEADER_SIZE_LEN);
	unsigned char* pheaddata = (unsigned char*)malloc(uheadlen);

	fread(pheaddata+ ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN, 1, uheadlen - ASF_HEADER_GUID_LEN - ASF_HEADER_SIZE_LEN, pSrcFile);
	unsigned char* pdecheaddata = (unsigned char*)malloc(uheadlen);
	int ioutlen = uheadlen - ASF_HEADER_GUID_LEN - ASF_HEADER_SIZE_LEN ;
	SM4DecECBMode(strkey, pheaddata + ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN, uheadlen - ASF_HEADER_GUID_LEN - ASF_HEADER_SIZE_LEN, pdecheaddata + ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN, &ioutlen);
	construct_asf_headsize(strheadlen,ASF_HEADER_SIZE_LEN,ioutlen + ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN);
	memcpy(pdecheaddata, strasfheader, ASF_HEADER_GUID_LEN);
	memcpy(pdecheaddata + ASF_HEADER_GUID_LEN, strheadlen, ASF_HEADER_SIZE_LEN);
	*pdechead = pdecheaddata;
	*udecheadlen = ioutlen + ASF_HEADER_GUID_LEN + ASF_HEADER_SIZE_LEN;
	*ioffset = (irvcheadlen + uheadlen - (*udecheadlen));
	safe_free(pheaddata);
	
	fclose(pSrcFile);
	iret = 0;
	return iret;
}


int rvc_free_data(void** pdechead)
{
	int iret = -1;

	safe_free(*pdechead);

	iret = 0;

	return iret;
}