// AESEncrypt.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "stdio.h"
#include "io.h"
#include "aes.h"
#include <string>
#include <vector>

extern "C"
{
	extern int encfile(FILE *fin, FILE *fout, aes_encrypt_ctx ctx[1], const char* ifn, const char* ofn);
	extern int decfile(FILE *fin, FILE *fout, aes_decrypt_ctx ctx[1], const char* ifn, const char* ofn);
};

// src and trg can be the same pointer (in place encryption)
// len must be in bytes and must be multiple of 8 byts (64bits).
// key is 128bit:  int key[4] = {n1,n2,n3,n4};
// void encipher(unsigned int *const v,unsigned int *const w,const unsigned int *const k )
#define TEA_ENCODE( src,trg,len,key ) {\
	register unsigned int *v = (src), *w = (trg), *k = (key), nlen = (len) >> 3; \
	register unsigned int delta=0x9E3779B9,a=k[0],b=k[1],c=k[2],d=k[3]; \
	while (nlen--) {\
	register unsigned int y=v[0],z=v[1],n=32,sum=0; \
	while(n-->0) { sum += delta; y += (z << 4)+a ^ z+sum ^ (z >> 5)+b; z += (y << 4)+c ^ y+sum ^ (y >> 5)+d; } \
	w[0]=y; w[1]=z; v+=2,w+=2; }}

// src and trg can be the same pointer (in place decryption)
// len must be in bytes and must be multiple of 8 byts (64bits).
// key is 128bit: int key[4] = {n1,n2,n3,n4};
// void decipher(unsigned int *const v,unsigned int *const w,const unsigned int *const k)
#define TEA_DECODE( src,trg,len,key ) {\
	register unsigned int *v = (src), *w = (trg), *k = (key), nlen = (len) >> 3; \
	register unsigned int delta=0x9E3779B9,a=k[0],b=k[1],c=k[2],d=k[3]; \
	while (nlen--) { \
	register unsigned int y=v[0],z=v[1],sum=0xC6EF3720,n=32; \
	while(n-->0) { z -= (y << 4)+c ^ y+sum ^ (y >> 5)+d; y -= (z << 4)+a ^ z+sum ^ (z >> 5)+b; sum -= delta; } \
	w[0]=y; w[1]=z; v+=2,w+=2; }}

// encode size ignore last 3 bits of size in bytes. (encode by 8bytes min)
#define TEA_GETSIZE( len ) ((len) & (~7))

//////////////////////////////////////////////////////////////////////////
int encfile_tea(FILE *fin, FILE *fout, aes_encrypt_ctx ctx[1], const char* ifn, const char* ofn)
{
	unsigned int key[4] = {76378631,637836237,115276122,1813576527};
	fseek(fin,0,SEEK_END);
	int size = ftell(fin);
	fseek(fin,0,SEEK_SET);
	char *data = (char*)malloc(size+10);
	fread( data,size,1,fin );
	TEA_ENCODE( (unsigned int*)data,(unsigned int*)data,TEA_GETSIZE(size),key );
	fwrite( data,size,1,fout );
	free(data);
	return 0;
}

//////////////////////////////////////////////////////////////////////////
int decfile_tea(FILE *fin, FILE *fout, aes_decrypt_ctx ctx[1], const char* ifn, const char* ofn)
{
	unsigned int key[4] = {76378631,637836237,115276122,1813576527};
	fseek(fin,0,SEEK_END);
	int size = ftell(fin);
	fseek(fin,0,SEEK_SET);
	char *data = (char*)malloc(size+10);
	fread( data,size,1,fin );
	TEA_DECODE( (unsigned int*)data,(unsigned int*)data,TEA_GETSIZE(size),key );
	fwrite( data,size,1,fout );
	free(data);
	return 0;
}


struct SScanInfo
{
	bool bEncode;
	aes_encrypt_ctx ectx[1];
	aes_decrypt_ctx dctx[1];
};

//////////////////////////////////////////////////////////////////////////
void EncryptFile( const char *infile,SScanInfo &si )
{
	// Encrypt file.
	char outfile[_MAX_PATH];
	char drive[_MAX_DRIVE];
	char dir[_MAX_DIR];
	char fname[_MAX_FNAME];
	char ext[_MAX_EXT];
	char outext[_MAX_EXT];
	_splitpath( infile,drive,dir,fname,ext );
	if (si.bEncode)
	{
		if (strstr(ext,"_e"))
			return;
		strcpy( outext,ext );
		strcat( outext,"_e" );
	}
	else
	{
		// Remove _e.
		strcpy( outext,ext );
		char *p = strstr(outext,"_e");
		if (!p)
			return;
		*p = 0;
	}
	_makepath( outfile,drive,dir,fname,outext );
	FILE *fin = fopen( infile,"rb" );
	FILE *fout = fopen( outfile,"wb" );
	if (si.bEncode)
	{
		printf( "Encrypting %s -> %s\n",infile,outfile );
		encfile(fin, fout, si.ectx, infile,outfile);
		//encfile_tea(fin, fout, si.ectx, infile,outfile);
	}
	else
	{
		printf( "Decrypting %s -> %s\n",infile,outfile );
		decfile(fin, fout, si.dctx, infile,outfile);
		//decfile_tea(fin, fout, si.dctx, infile,outfile);
	}
	fclose(fin);
	fclose(fout);
}

//////////////////////////////////////////////////////////////////////////
// Enumerate folder.
void ScanFolder( const char *sFolder,SScanInfo &si )
{
	_finddata_t fd;

	std::vector<std::string> files;
	char sSearch[_MAX_PATH];
	sprintf( sSearch,"%s\\*.*",sFolder );
	intptr_t handle = _findfirst( sSearch,&fd );
	if (handle != -1)
	{
		do
		{
			if (fd.attrib & _A_SUBDIR)
			{
				if (strcmp(fd.name,".") != 0 && strcmp(fd.name,"..") != 0)
				{
					char sPath[_MAX_PATH];
					sprintf( sPath,"%s\\%s",sFolder,fd.name );
					ScanFolder( sPath,si );
				}
			}
			else
			{
				char infile[_MAX_PATH];
				sprintf( infile,"%s\\%s",sFolder,fd.name );
				files.push_back(infile);
			}
		}
		while (_findnext(handle,&fd) == 0);
		_findclose(handle);
	}

	for (int i = 0; i < (int)files.size(); i++)
	{
		EncryptFile( files[i].c_str(),si );
	}
}

int main(int argc, char* argv[])
{
	char key[32] = "3CE2698701289029F3926DF0191189A";

	if (argc < 2)
	{
		// Usage
		printf( "Usage:\n" );
		printf( "  AESEncrypt directory [KeyFileName]\n" );
		printf( "  To recursively encrypt all files in the specified directory.\n" );
		printf( "  KeyFileName - Optional filename of binary key file, first 32 bytes of this file will be used as a key.\n" );
		printf( "\n" );
		printf( "\n" );
		return 0;
	}

	if (argc > 2)
	{
		// Read key file.
		FILE *file = fopen( argv[2],"rb" );
		if (file)
		{
			if (fread( key,1,32,file ) != 32)
			{
				printf( "Key file %s must be at least 32 bytes long.",argv[2] );
				fclose(file);
				return 0;
			}
			fclose(file);
		}
		else
		{
			printf( "Key file %s not found.",argv[2] );
			return 0;
		}
	}

	SScanInfo si;
	si.bEncode = true;
	aes_encrypt_key((unsigned char*)key, 32, si.ectx);
	aes_decrypt_key((unsigned char*)key, 32, si.dctx);

	ScanFolder( argv[1],si );

	return 0;
}