#include <iostream>
#include <string>

using namespace std;

#include "cryptlib.h"
#include "salsa.h"
#include "panama.h"
#include "aes.h"
#include "modes.h"
#include "osrng.h"
#include "secblock.h"
#include "filters.h"
#include "hex.h"

using namespace CryptoPP;

typedef std::auto_ptr<SymmetricCipher> SymmetricCipherPtr;

int main(int argc, char** argv)
{
	string streamCipher("salsa");
	SymmetricCipherPtr encryptor;
	SymmetricCipherPtr decryptor;
	string specifiedCipher;
	if(argc > 1) {
		specifiedCipher = string(argv[1]);
	} else {
		specifiedCipher = streamCipher;
		cerr << "No cipher was specified. Using " << streamCipher << endl;
	}
	unsigned int keyLen = 0;
	unsigned int ivLen = 0;
	if(specifiedCipher == "salsa") {
		SymmetricCipherPtr et(new Salsa20::Encryption());
		SymmetricCipherPtr dt(new Salsa20::Decryption());
		encryptor = et;
		decryptor = dt;
		keyLen = Salsa20::MAX_KEYLENGTH;
		ivLen = Salsa20::IV_LENGTH;
	} else if(specifiedCipher == "panama") {
		SymmetricCipherPtr et(new PanamaCipher<LittleEndian>::Encryption());
		SymmetricCipherPtr dt(new PanamaCipher<LittleEndian>::Decryption());
		encryptor = et;
		decryptor = dt;
		keyLen = PanamaCipher<LittleEndian>::MAX_KEYLENGTH;
		ivLen = PanamaCipher<LittleEndian>::IV_LENGTH;
	} else if(specifiedCipher == "aes128cbc") {
		SymmetricCipherPtr et(new CBC_Mode<AES>::Encryption() );
		SymmetricCipherPtr dt(new CBC_Mode<AES>::Decryption() );
		encryptor = et;
		decryptor = dt;
		keyLen = 16;
		ivLen = AES::BLOCKSIZE;
	} else if(specifiedCipher == "aes256cbc") {
		SymmetricCipherPtr et(new CBC_Mode<AES>::Encryption() );
		SymmetricCipherPtr dt(new CBC_Mode<AES>::Decryption() );
		encryptor = et;
		decryptor = dt;
		keyLen = 32;
		ivLen = AES::BLOCKSIZE;
	} else {
		cerr << "An unrecognized cipher was specified." << endl;
		return 1;
	}

	AutoSeededRandomPool rng;
	SecByteBlock key(keyLen);
	rng.GenerateBlock(key,keyLen);
	SecByteBlock iv(ivLen);
	
	// !!! NOTE: This may not really be appropriate !!!
	rng.GenerateBlock(iv,ivLen);
	
	
	string plain("Some super secret text...");
	string cipher;
	
	if(ivLen) {
		encryptor->SetKeyWithIV(key,keyLen,iv);
	}else{
		encryptor->SetKey(key,keyLen);
	}
	
	StringSource ps(plain,true,new StreamTransformationFilter(*encryptor,new HexEncoder(new StringSink(cipher))));
	cout << "Plaintext: " << plain << endl;
	cout << "Ciphertext: " << cipher << endl;
	
	string recovered;
	
	if(ivLen) {
		decryptor->SetKeyWithIV(key,keyLen,iv);
	} else {
		decryptor->SetKey(key,keyLen);
	}
	StringSource rs(cipher,true,new HexDecoder(new StreamTransformationFilter(*decryptor,new StringSink(recovered))));
	cout << "Recovered: " << recovered << endl;
	
	return 0;
	
}