#include "stdafx.h"
#include "filecryption.h"
#include "fileanalysis.h"
#include "rvcfileheader.h"
#include "asf.h"
#include "mp4info.h"
#include <string.h>

#include <iostream>
using namespace std;

#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)  \
		do{\
			if(x) { free(x); x = NULL; }\
		}while(0)
#endif



static int encrypt_asf_file(char* poutfilename, char* phead, uint32_t uheaderlen, unsigned char* pstrkey, uint32_t ukeylen, FILE * pSrcFile, const filecryption_callback_t * pcallback)
{
	int iret = -1;

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

	//1. 新增自定义文件头
	fwrite(phead, 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, FILECRYPT_LEVEL_ERROR, "%s", "math guid failed, current format is not surpport!");
		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(pstrkey, 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(pDestFile);

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

	iret = 0;

	return iret;
}


static int encrypt_mp4_file(char* poutfilename, char* phead, uint32_t uheaderlen, unsigned char* pstrkey, uint32_t ukeylen, FILE* pSrcFile, const filecryption_callback_t* pcallback)
{
	int iret = -1;

	unsigned int unoencryptlen = get_noencrypt_boxs_size(pSrcFile, pcallback, true);
	if (0 == unoencryptlen) {
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "get not encrypt boxs size failed.");
		return iret;
	}
	//safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "not encrypt file length is %u.", unoencryptlen);

	FILE* pDestFile = fopen(poutfilename, "wb");
	if (NULL == pDestFile) {
		fclose(pSrcFile);
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "open output file failed.");
		return iret;
	}

	fseek(pSrcFile, 0, SEEK_END);
	uint32_t usrcfilelen = (uint32_t)ftell(pSrcFile);
	rewind(pSrcFile);
	safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "src file length is %u.", usrcfilelen);

	//1. 新增自定义文件头
	fwrite(phead, 1, uheaderlen, pDestFile);

	//2. 填充ftyp, free, mdat box
	char* data_buffer = (char*)malloc(RVC_READ_BUFFER_SIZE);
	unsigned int uleftfilelen = unoencryptlen;
	do {
		if (RVC_READ_BUFFER_SIZE == fread(data_buffer, 1, RVC_READ_BUFFER_SIZE, pSrcFile)) {
			fwrite(data_buffer, 1, RVC_READ_BUFFER_SIZE, pDestFile);
		}
		uleftfilelen -= RVC_READ_BUFFER_SIZE;
	} while (uleftfilelen >= RVC_READ_BUFFER_SIZE);

	if (uleftfilelen > 0) {
		if (uleftfilelen == fread(data_buffer, 1, uleftfilelen, pSrcFile)) {
			fwrite(data_buffer, 1, uleftfilelen, pDestFile);
		}
	}

	//3. 读取moov box信息
	//size_t udatalen = 0;
	//while (udatalen = fread(data_buffer, 1, RVC_READ_BUFFER_SIZE, pSrcFile)) {
	//	fwrite(data_buffer, 1, udatalen, pDestFile);
	//}

	uint32_t umoovboxlen = usrcfilelen - unoencryptlen;
	unsigned char* pmoovboxdata = (unsigned char*)malloc(umoovboxlen);
	fread(pmoovboxdata, 1, umoovboxlen, pSrcFile);

	int32_t iencdatalen = umoovboxlen + SM4ENC_BLOCK_SIZE;
	unsigned char* pobjdataenc = (unsigned char*)malloc(iencdatalen);
	
	if (0 == SM4EncECBMode(pstrkey, pmoovboxdata, umoovboxlen, pobjdataenc, &iencdatalen)) {
		safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "after encode, moov box length is %u.", iencdatalen);
		if (iencdatalen > 0) {
			//3. 填充加密后的moov box长度
			fwrite(&iencdatalen, 1, sizeof(uint32_t), pDestFile);
			//4. 填充加密后的moov box内容
			fwrite(pobjdataenc, 1, iencdatalen, pDestFile);
			iret = 0;
		}
	}
	
	if(0 != iret){
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "SM4 encrypt moovbox failed, and src moovbox len is %d.", umoovboxlen);
	}
	
	fclose(pDestFile);

	safe_free(data_buffer);
	safe_free(pmoovboxdata);
	safe_free(pobjdataenc);

	return iret;
}


int encryption_file(char* poutfile, uint32_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};
	uint32_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};
	uint32_t uindex = (RVC_FILE_HEADER_FLAG_LEN + sizeof(uint32_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;
	}

	if (eVerA == eversion) {
		iret = encrypt_asf_file(poutfile, pbuffer, uheaderlen, strkey, SM4ENC_BLOCK_SIZE, pSrcFile, pcallback);
	}
	else {
		iret = encrypt_mp4_file(poutfile, pbuffer, uheaderlen, strkey, SM4ENC_BLOCK_SIZE, pSrcFile, pcallback);
	}

	fclose(pSrcFile);
	
	return iret;
}


static int decrypt_asf_file(char* poutfilename, FILE* pSrcFile, unsigned char* pstrkey, uint32_t ukeylen, const filecryption_callback_t* pcallback)
{
	int iret = -1;
	FILE* pDestFile = fopen(poutfilename, "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(pstrkey, 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(pDestFile);

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

	return iret;
}


static int decrypt_mp4_file(char* poutfilename, FILE* pSrcFile, unsigned char* pstrkey, uint32_t ukeylen, const filecryption_callback_t* pcallback)
{
	int iret = -1;
	FILE* pDestFile = fopen(poutfilename, "wb");
	if (NULL == pDestFile) {
		return iret;
	}

	unsigned int unoencryptlen = get_noencrypt_boxs_size(pSrcFile, pcallback, false);
	if (0 == unoencryptlen) {
		return iret;
	}
	safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "not encrypt file length is %u.", unoencryptlen);

	//1. 填充ftyp, free, mdat box
	char* data_buffer = (char*)malloc(RVC_READ_BUFFER_SIZE);
	unsigned int uleftfilelen = unoencryptlen;
	do {
		if (RVC_READ_BUFFER_SIZE == fread(data_buffer, 1, RVC_READ_BUFFER_SIZE, pSrcFile)) {
			fwrite(data_buffer, 1, RVC_READ_BUFFER_SIZE, pDestFile);
		}
		uleftfilelen -= RVC_READ_BUFFER_SIZE;
	} while (uleftfilelen >= RVC_READ_BUFFER_SIZE);

	if (uleftfilelen > 0) {
		if (uleftfilelen == fread(data_buffer, 1, uleftfilelen, pSrcFile)) {
			fwrite(data_buffer, 1, uleftfilelen, pDestFile);
		}
	}

	//2. 获取加密后的moov box大小
	uint32_t uencryptboxsize = 0;
	fread(&uencryptboxsize, 1, sizeof(uint32_t), pSrcFile);
	safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "encrypt moov box size is = %d.", uencryptboxsize);

	if (uencryptboxsize > 0) {
		//3. 读取加密后的moov box
		unsigned char* pmoovboxdata = (unsigned char*)malloc(uencryptboxsize);
		unsigned char* pdecmoovboxdata = (unsigned char*)malloc(uencryptboxsize);
		int ioutlen = uencryptboxsize;
		fread(pmoovboxdata, 1, uencryptboxsize, pSrcFile);

		//4. 解密moov box
		SM4DecECBMode(pstrkey, pmoovboxdata, uencryptboxsize, pdecmoovboxdata, &ioutlen);
		safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "decrypt moov box size is = %d.", ioutlen);

		//5. 填充解密后的moov box
		fwrite(pdecmoovboxdata, 1, ioutlen, pDestFile);

		safe_free(pdecmoovboxdata);
		safe_free(pmoovboxdata);

		iret = 0;
	}
	else {
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "get encrypt moovbox failed, decrypt failed!");
	}

	fclose(pDestFile);
	safe_free(data_buffer);

	return iret;
}

int decryption_file(char* poutfile, uint32_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, FILECRYPT_LEVEL_ERROR, "%s","invalid encryption version param.");
		return iret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "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, FILECRYPT_LEVEL_INFO, "file %s is not encryption!",pfilename);
		fclose(pSrcFile);
		return iret;
	}

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

	if (0 != memcmp(prvcbuffer+sizeof(uint32_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN)){
		if (0 == memcmp(prvcbuffer+sizeof(uint32_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN-1)){
			safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "file %s is encrypt but decrypt version is not matched!",pfilename);
		} 
		fclose(pSrcFile);
		safe_free(prvcbuffer);
		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(uint32_t), pcallback)){
		safe_free(prvcbuffer);
		return iret;
	}

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

	if (eVerA == eversion) {
		iret = decrypt_asf_file(poutfile, pSrcFile, strkey, SM4ENC_BLOCK_SIZE, pcallback);
	}
	else {
		iret = decrypt_mp4_file(poutfile, pSrcFile, strkey, SM4ENC_BLOCK_SIZE, pcallback);
	}

	fclose(pSrcFile);
	safe_free(prvcbuffer);

	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, FILECRYPT_LEVEL_ERROR, "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 == strncmp(strrvcheader, rvc_header, RVC_FILE_HEADER_FLAG_LEN)){
			bret = true;
		} 
		else{
			safe_log(pcallback, FILECRYPT_LEVEL_INFO, "file %s is not encryption!",pfilename);
		}
	}
	fclose(pSrcFile);

	return bret;
}


/*
文件解密功能；
	输入：文件路径
	输出：文件明文头数据，文件明文头长度，加解密前后文件偏移量，自定义头文件头信息
*/
int rvc_file_decrypt(unsigned char** pdechead, uint32_t* udecheadlen, int* ioffset, char** pstrjson, uint32_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, FILECRYPT_LEVEL_ERROR, "%s","invalid cryption version param.");
		return iret;
	}

	FILE* pSrcFile = fopen(pfilename, "rb");
	if (NULL == pSrcFile){
		safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "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, FILECRYPT_LEVEL_INFO, "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(uint32_t), pSrcFile);
		char* prvcbuffer = (char*)malloc(irvcheadlen);
		memset(prvcbuffer, 0, irvcheadlen);
		fread(prvcbuffer, 1, irvcheadlen - RVC_FILE_HEADER_FLAG_LEN - sizeof(uint32_t), pSrcFile);

		if (0 != memcmp(prvcbuffer+sizeof(uint32_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN)){
			if (0 == memcmp(prvcbuffer+sizeof(uint32_t), cryption_ver_flag_table[eversion], RVC_CRYPTION_VER_FLAG_LEN-1)){
				safe_log(pcallback, FILECRYPT_LEVEL_ERROR, "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(uint32_t), 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(uint32_t));
		if (0 == ijson){
			safe_log(pcallback, FILECRYPT_LEVEL_DEBUG, "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, FILECRYPT_LEVEL_ERROR, "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;
}


bool is_file_completed(const char* pfilename, const filecryption_callback_t* pcallback)
{
	return is_mp4file_completed(pfilename, pcallback);
}

int get_file_sm3digest(char* strbuf, uint32_t ubuflen, const char* pfilename)
{
	return get_file_hashstring(strbuf, ubuflen, pfilename);
}