#include <fstream>
#include <stdexcept>
#include <crypto++/sha.h>
#include <crypto++/modes.h>
#include <crypto++/aes.h>

using namespace std;
using namespace CryptoPP;

void encrypt_ecb(char *data, const int size,
				 const string& key, const string& file)
{
	byte hash[20];
	SHA1 sha1;
	sha1.CalculateDigest(hash,(byte*)&key[0],key.size());
	ECB_Mode<AES>::Encryption aes(hash,AES::DEFAULT_KEYLENGTH);

	//NOTE: if(size % aes.MandatoryBlockSize()!=0) up to
	//aes.MandatoryBlockSize()-1 trailing bytes from the file will
	//NOT be encrypted when using ECB mode!
	aes.ProcessData((byte*)data,(byte*)data,size);
	ofstream out(file.c_str(),ios::binary);
	out.write(data,size);
}

void encrypt_ctr(char *data, const int size,
				 const string& key, const string& file)
{
	byte hash[20];
	SHA1 sha1;
	sha1.CalculateDigest(hash,(byte*)&key[0],key.size());
	CTR_Mode<AES>::Encryption aes(hash,AES::DEFAULT_KEYLENGTH,hash);

	aes.ProcessData((byte*)data,(byte*)data,size);
	ofstream out(file.c_str(),ios::binary);
	out.write(data,size);
}

int main(int argc, char *argv[])
{
	if(argc!=4) throw runtime_error("usage: encrypt file key mode");
	string file(argv[1]);
	string key(argv[2]);
	string mode(argv[3]);

	ifstream in(file.c_str(),ios::binary);
	in.seekg(0,ios::end);
	int size=in.tellg();
	in.seekg(0,ios::beg);
	char *data=new char[size];
	in.read(data,size);

	file+=".enc";
	if(mode=="ecb") encrypt_ecb(data,size,key,file);
	else if(mode=="ctr") encrypt_ctr(data,size,key,file);
	else throw runtime_error("Mode is either ecb or ctr");
	delete[] data;
}