#include<string.h>
#include<stdio.h>
#include<stdlib.h>
#include "rijndael.h"

void display(UINT8 *buffer, int len, int mode = 0);
void main (int argc, char * argv[])
{
	enum Format{txt, jpg};        //運(yùn)行模式：ascii文件，jpg圖形
	Format format = txt;
	UINT8 key[16];                //key值
	UINT8 *plainText;             //明文所在的緩沖區(qū)
	UINT8 *output;                //密文所在的緩沖區(qū)
	UINT8 *output2;               //brute force時(shí)的臨時(shí)密文緩沖區(qū) 
	char buffer[150];
	char *fn=NULL;
	bool strict;                  //jpg模式下是否采用嚴(yán)格標(biāo)準(zhǔn)的jpg文件格式頭
	bool decrypted = false;			//為true時(shí)，對(duì)輸入進(jìn)行破解，此時(shí)，參數(shù)-b無(wú)效
	bool bruteforce = true;        //加密之后是否進(jìn)行brute force
	FILE *pT = NULL;               //指向明文所在文件
	FILE *pK = NULL;				//指向Key所在文件
	FILE *pE = NULL;			   //指向密文所在文件
	FILE *pR = NULL;               //記錄運(yùn)行情況的文件
	FILE *pJ = NULL;               //jpg模式下的.jpg文件
	int len;                       //密文長(zhǎng)度
	int len1;                      //burte force時(shí)的臨時(shí)密文長(zhǎng)度
	int keyLen = 0;				   //key的長(zhǎng)度(byte)
	int plainTextLen;			   //明文的長(zhǎng)度
	int i;
	
	//處理輸入?yún)?shù)
	i = 1;

	//第一個(gè)參數(shù)可為不帶'-'的明文文件名,只有當(dāng)文件后綴為.jpg時(shí)，才將模式設(shè)置為jpg，其他情況一概視為txt
	if(argc > 1 && argv[i][0] != '-')
	{
		if(strlen(argv[1]) > strlen(".jpg") && !strcmp(&argv[1][strlen(argv[1])-strlen(".jpg")], ".jpg"))
			format = jpg;
		if((pT = fopen( argv[i++], "rb" )) == NULL)
		{
			printf( "The file %s could not be opened for read purpose\n", argv[1]);
			exit(1);
		}
	}

	/*接下去分析其他輸入?yún)?shù)，其中
	1. -e encrypted已被加密的文件
	2. -t plaintext.txt指定明文文件，并設(shè)置為txt模式
	3. -j pic.jpg 指定明文文件，并設(shè)置為jpg模式
	4. -k key指定key值，其由不多于16個(gè)0-9組成
	5. -b 表示不進(jìn)行brute force
	6. -o 指定加密后的輸出文件名
	*/
	while(i < argc)	
	{
		if(argv[i][0] == '-')
		{
			if(strlen(argv[i]) < 2)
			{
				printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
				exit(1);
			}
			switch(argv[i][1])
			{
			case 'e':
				decrypted = true;
				if(i == argc -1)
				{
					printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
					exit(1);
				}
				if(pT)
				{
					printf("The input file has been designted more than once\n");
					exit(1);
				}
				if((pT = fopen( argv[++i], "rb" )) == NULL )
				{
					printf( "The file %s could not be opened for read purpose\n", argv[i]);
					exit(1);
				}
				break;
			case 't':
			case 'j':

				if(argv[i][1] == 'j')
					format = jpg;
				else format = txt;
			
				if(i == argc -1)
				{
					printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
					exit(1);
				}
				if(pT)
				{
					printf("The plaintext file has been designted more than once\n");
					exit(1);
				}
				if((pT = fopen( argv[++i], "rb" )) == NULL )
				{
					printf( "The file %s could not be opened for read purpose\n", argv[i]);
					exit(1);
				}
				break;
			case 'k':
				if(i == argc -1)
				{
					printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
					exit(1);
				}
				i++;
				memset(key, 0, 16);				
				keyLen = strlen(argv[i]) < 16? strlen(argv[i]):16;
				memcpy(key, argv[i], keyLen);
				break;
			case 'o':
				if(i == argc -1)
				{
					printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
					exit(1);
				}
				if(pE)
				{
					printf("The encrypted ouput file has been designted more than once\n");
					exit(1);
				}
				if((pE = fopen( argv[++i], "wb" )) == NULL )
				{
					fn=argv[i];
					printf( "The file %s could not be opened for write purpose\n", argv[i]);
					exit(1);
				}
				break;
			case 'b':
				bruteforce = false;
				break;
			default:
				printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
				exit(1);
				break;
			}
		}
		else 
		{
			printf("Usage: rijndeal [-e encrypted | -t plaintext.txt | -j pic.jpg ] [-k key] [-o encrypted] [-b]\n");
			exit(1);
		}
		i++;
	}
	
	//如未指定明文文件，采用默認(rèn)的當(dāng)前目錄下的plaintext.txt
	if(pT == NULL)
		if((pT = fopen( "plaintext.txt", "rb" )) == NULL )
		{
			printf( "The file 'plainText.txt' could not be opened for read purpose\n");
			exit(1);
		}
		
	//如未指定key值，則從默認(rèn)的當(dāng)前目錄下文件key讀出
	if(keyLen == 0)
	{
		memset(key, 0, 16);
		if((pK = fopen( "key", "rb")) == NULL )
		{
			printf("The file 'key' could not be opened\n");
			exit(1);
		}
		keyLen = fread(key, sizeof(UINT8), 16, pK);
		fclose(pK);
	}
    
	//運(yùn)行情況記錄在當(dāng)前目錄下的result文件中
	if((pR = fopen("result", "wb")) == NULL)
	{
		printf("The file 'result' couldn't be opend for write purpose\n");
		exit(1);
	}

	//讀取明文，保存至plainText中
	if(fseek(pT, 0, SEEK_END))
	{
		printf("The pointer of file %s couldn't be sought\n", argv[1]);
		exit(1);
	}

	if((plainTextLen = ftell(pT)) == -1)
	{
		printf("Error occur when getting the length of the file %s\n", argv[1]);
		exit(1);
	}

	if(fseek(pT, 0, SEEK_SET))
	{
		printf("The pointer of file %s couldn't be sought\n", argv[1]);
		exit(1);
	}
	
	plainText = new UINT8[plainTextLen];
	plainTextLen = fread(plainText, sizeof(UINT8), plainTextLen, pT);
	fclose(pT);

	//如在.jpg模式下，則判斷.jpg文件是否是標(biāo)準(zhǔn)的.jpg格式
	if( format == jpg && !decrypted )
	{
		strict = true;
		if(memcmp(plainText, "\xFF\xD8\xFF\xE0\x00\x10JFIF", 10))
		{
			strict = false;
			if(memcmp(plainText, "\xFF\xD8\xFF", 3))
			{
				printf("file %s isn't compitable with the jpg format\n", argv[1]);
				exit(1);
			}
		}
	}

	//輸出key值到控制臺(tái)和result文件中，記錄運(yùn)行狀況
	sprintf(buffer, "Key(%d):\t", 16);
	printf(buffer);
	display(key, 16, 1);
	fwrite(buffer, sizeof(char), strlen(buffer), pR);
	fwrite(key, sizeof(UINT8), 16, pR);

	sprintf(buffer, "\n-------------------------------------------------------------------\n");
	printf(buffer);
	fwrite(buffer, sizeof(char), strlen(buffer), pR);

	//在txt模式下，輸出明文到控制臺(tái)和result文件
	if(format == txt && !decrypted)
	{
		sprintf(buffer, "Plaintext(%d):\t", plainTextLen);
		printf(buffer);
		fwrite(buffer, sizeof(char), strlen(buffer), pR);

		display(plainText, plainTextLen, 1);
		fwrite(plainText, sizeof(UINT8), plainTextLen, pR);

		sprintf(buffer,"\n------------------------------------------------------------\n");
		printf(buffer);
		fwrite(buffer, sizeof(char), strlen(buffer), pR);
	
	}
///////////////////////////////////////////////////////////////////////////////////
///////////////////////      Decrypt       ///////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////
	if(decrypted)
	{
		Rijndael rin;
		rin.init(Rijndael::CBC, Rijndael::Decrypt, key, Rijndael::Key16Bytes);
		output = new UINT8[plainTextLen];
		len = rin.padDecrypt(plainText, plainTextLen, output);
		if(len >= 0)
		{			
			sprintf(buffer, "Decrypt(%d) succeeds\n",len);
			printf(buffer);
			fwrite(buffer, sizeof(char), strlen(buffer), pR);

			fclose(pR);
			if(pE == NULL)
				pE = fopen("Encrypted", "wb");
			if(pE)
			{
				fwrite(output, sizeof(UINT8), len, pE);	
				fclose(pE);
			}
		}
		else
		{
			printf("Encrypted data has been corrupted\n");
			exit(1);
		}
		exit(0);
	}
////////////////////////////////////////////////////////////////////////////////////
//////////////////////         Encrypt      ////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////
//加密明文，利用Rijndael的實(shí)例rin來(lái)實(shí)現(xiàn)
	output = new UINT8[plainTextLen + 16];
	
	Rijndael rin;
	
	rin.init(Rijndael::CBC, Rijndael::Encrypt,key, Rijndael::Key16Bytes);
	len = rin.padEncrypt(plainText, plainTextLen, output);

	if(format==txt)
		printf("text mode:\n\n");
	else
	{
		printf("jpg mode:\n");
		printf("As the jpg files is always big, the possible result is saved in the corresponding files but not displayed here!\n\n");
	}
	if(len >= 0)
	{
		printf("Encrypted(%d):\t\n",len);
		if(format == txt)
		{
			display(output, len);	
			printf("\n");
		}
		if(fn==NULL)
			printf("The result is saved in the file \"Encrypted\"!\n");
		else
			printf("The result is saved in the file \"%s\"!\n",fn);
		
		printf("\n------------------------------------------------------------\n");
		
		if(pE == NULL)
			pE = fopen("Encrypted", "wb");
		if(pE)
		{
			fwrite(output, sizeof(UINT8), len, pE);	
			fclose(pE);
		}
	}
	else
	{
		printf("Error occur during encrypting\n");
		exit(1);
	}
	if(!bruteforce) 
		exit(0);
///////////////////////////////////////////////////////////////////////////////////////////////////	
/////////////////         try bruteforce                 //////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////
//記錄運(yùn)行狀況
	//sprintf(buffer,"------------------------------------------------------------\n");
//	printf(buffer);
	fwrite(buffer, sizeof(char), strlen(buffer), pR);
	sprintf(buffer,"--------------------Try to bruteforce-----------------------\n");
	printf(buffer);
	fwrite(buffer, sizeof(char), strlen(buffer), pR);
	sprintf(buffer,"------------------------------------------------------------\n");
	printf(buffer);
	fwrite(buffer, sizeof(char), strlen(buffer), pR);

	output2 = new UINT8[len];
	for( i = 0; i < keyLen; i++)
	{
		key[i] = '0';
	}
//txt模式下下的brute force
	if(format == txt)
	{
		//printf("Txt mode:\n");
		while(1)
		{
			rin.init(Rijndael::CBC,Rijndael::Decrypt,key,Rijndael::Key16Bytes);
			len1 = rin.padDecrypt(output,len,output2);
		
			if(len1 >= 0)
			{
				for(i = 0; i < len1 && !(output2[i]&0x80); i++);
				if(i == len1) 
				{
					sprintf(buffer, "Key(%d):\t", 16);
					printf(buffer);
					display(key, 16, 1);
					fwrite(buffer, sizeof(char), strlen(buffer), pR);
					fwrite(key, sizeof(UINT8), 16, pR);
				
					sprintf(buffer, "\nMessage(%d)\t", len1);
					printf(buffer);
					display(output2, len1, 1);

					fwrite(buffer, sizeof(char), strlen(buffer), pR);
					fwrite(output2, sizeof(UINT8), len1, pR);
					sprintf(buffer, "\n--------------------------------------------------------------------\n");
					printf(buffer);
					fwrite(buffer, sizeof(char), strlen(buffer), pR);
				}
			}
			for(i = 0; i < keyLen; i++)
			{
				key[i] = '0'+(++key[i] - '0') % 10;
				if(i == (keyLen - 1) && key[i] == '0')
				{
					printf("The decrypted data saved in the file \"result\"!!\n\n");
					printf("Try bruteforce end\n");
					exit(1);
				}
				if(key[i]-'0')
					break;				
			}       
		}	
	}
//jpg模式下brute force
	else
	{
	//	printf("jpg mode:\n");
		fn=new char[150];
		while(1)
		{
			rin.init(Rijndael::CBC,Rijndael::Decrypt,key,Rijndael::Key16Bytes);
			//在此先破解第一塊，看是否符合jpg文件的頭部格式，如不符合，則無(wú)需破解其他明文塊。
			len1 = rin.blockDecrypt(output, 128, output2);

			if((strict && !memcmp(output2, "\xFF\xD8\xFF\xE0\x00\x10JFIF", 10) )|| (!strict && !memcmp(output2, "\xFF\xD8\xFF", 3)))
			{
				rin.init(Rijndael::CBC,Rijndael::Decrypt,key,Rijndael::Key16Bytes);
				len1 = rin.padDecrypt(output,len,output2);
				if(len1>=0)
				{
					memcpy(fn, key, keyLen);
					memcpy(fn+keyLen, ".jpg", 5);
					if((pJ = fopen(fn, "wb")) == NULL)
					{
						printf("file %s cann't be created", fn);
						exit(3);
					}
					fwrite(output2, sizeof(UINT8), len1, pJ);
					fclose(pJ);
					

					sprintf(buffer, "\nFind key:");
					printf(buffer);
					fwrite(buffer, sizeof(char), strlen(buffer), pR);

					display(key,16);
					fwrite(key, sizeof(UINT8), 16, pR);	
					printf("\nThe decrypted data saved in the file \"%s\"!!\n",fn);

					sprintf(buffer, "\n--------------------------------------------------------------------\n");
					printf(buffer);
					fwrite(buffer, sizeof(char), strlen(buffer), pR);
				}
			}
		//遍歷所有可能的key值
				
			for(i = 0; i < keyLen; i++)
			{
				key[i] = '0' + (++key[i] - '0') % 10;
				if(i == (keyLen - 1) && key[i] == '0')
				{
					printf("Try burteforce end\n");
					exit(1);
				}
				if(key[i]-'0')
					break;	
			}
		}
	
	}
	fclose(pR);
}

//輸出信息到控制臺(tái)
void display(UINT8 *buffer, int len, int mode)
{
	int i;

	if(mode < 2)
		for(i = 0; i < len; i++)
		{
			printf("%x%x ", buffer[i] >> 4, buffer[i] & 0x0f);
			if(i == len - 1) 
				break;
			if(i%16 == 15)
				printf("\n\t\t");
			else if(i%4 == 3) 
				printf("\t");
		}

	if(mode)
	{
		printf("\nASCII format(%d):", len);
		if(len > 16 )printf("\n");
		for(i = 0; i < len; i++)
			printf("%c",buffer[i]);

	}
}