#include <iostream>
using std::cout;
using std::cerr;
using std::endl;

#include <fstream>
using std::ofstream;

#include <string>
using std::string;

#include "hex.h"
using CryptoPP::HexEncoder;

#include "cryptlib.h"
using CryptoPP::byte;
using CryptoPP::BufferedTransformation;
using CryptoPP::AuthenticatedSymmetricCipher;

#include "secblock.h"
using CryptoPP::SecByteBlock;

#include "filters.h"
using CryptoPP::StringSink;
using CryptoPP::StringSource;
using CryptoPP::AuthenticatedEncryptionFilter;
using CryptoPP::AuthenticatedDecryptionFilter;
using CryptoPP::Redirector;

#include "files.h"
using CryptoPP::FileSource;

#include "aes.h"
using CryptoPP::AES;

#include "eax.h"
using CryptoPP::EAX;

class AuthenticatedDecryptionFilterBugFixed: public CryptoPP::AuthenticatedDecryptionFilter
{
public:
	AuthenticatedDecryptionFilterBugFixed(CryptoPP::AuthenticatedSymmetricCipher &c,
		CryptoPP::BufferedTransformation *attachment = NULLPTR,
		CryptoPP::word32 flags = DEFAULT_FLAGS,
		int truncatedDigestSize=-1,
		BlockPaddingScheme padding = DEFAULT_PADDING)
		: AuthenticatedDecryptionFilter(c, attachment, flags, truncatedDigestSize, padding) {}

	size_t AuthenticatedDecryptionFilterBugFixed::ChannelPutModifiable2(const std::string &channel,
		CryptoPP::byte *begin, size_t length, int messageEnd, bool blocking)
	{
		if (channel.empty())
		{
			if (m_lastSize > 0)
				m_hashVerifier.ForceNextPut();
			return FilterWithBufferedInput::PutModifiable2(begin, length, messageEnd, blocking);
		}
		if (channel == CryptoPP::AAD_CHANNEL)
			return m_hashVerifier.PutModifiable2(begin, length, 0, blocking);

		throw InvalidChannelName("AuthenticatedDecryptionFilter", channel);
	}
};

class AuthenticatedDecryptionFilterBugFixed2: public CryptoPP::AuthenticatedDecryptionFilter
{
public:
	AuthenticatedDecryptionFilterBugFixed2(CryptoPP::AuthenticatedSymmetricCipher &c,
		CryptoPP::BufferedTransformation *attachment = NULLPTR,
		CryptoPP::word32 flags = DEFAULT_FLAGS,
		int truncatedDigestSize=-1,
		BlockPaddingScheme padding = DEFAULT_PADDING)
		: AuthenticatedDecryptionFilter(c, attachment, flags, truncatedDigestSize, padding) {}

	size_t AuthenticatedDecryptionFilterBugFixed2::ChannelPutModifiable2(const std::string &channel,
		CryptoPP::byte *begin, size_t length, int messageEnd, bool blocking)
	{
		return ChannelPut2(channel, begin, length, messageEnd, blocking);
	}
};

int main(int argc, char* argv[])
{
	const int TAG_SIZE = 16;

	// Encrypted, with Tag
	string cipher, encoded;

	// Recovered
	string radata, rpdata;

	SecByteBlock key(32);
	memset( key, 0, key.size() );

	byte iv[12];
	memset( iv, 0, sizeof(iv) );

	string adata( 16, (char)0x00 );
	string pdata( 16, (char)0x00 );

	/*********************************\
		Section 1
	\*********************************/

	// Pretty print
	cout << "Crypto++ Version: " << CryptoPP::LibraryVersion() << endl << endl;

	encoded.clear();
	StringSource( key, key.size(), true, new HexEncoder( new StringSink( encoded ) ) );
	cout << "key: " << encoded << endl;

	encoded.clear();
	StringSource( iv, sizeof(iv), true, new HexEncoder( new StringSink( encoded ) ) );
	cout << "iv: " << encoded << endl;

	encoded.clear();
	StringSource( adata, true, new HexEncoder( new StringSink( encoded ) ) );
	cout << "adata: " << encoded << endl;

	encoded.clear();
	StringSource( pdata, true, new HexEncoder( new StringSink( encoded ) ) );
	cout << "pdata: " << encoded << endl;

	/*********************************\
		Section 2
	\*********************************/

	try
	{
		EAX< AES >::Encryption e;
		e.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		AuthenticatedEncryptionFilter ef( e, new StringSink( cipher ), false, TAG_SIZE );

		// AuthenticatedEncryptionFilter::ChannelPut
		//  defines two channels: "" (empty) and "AAD"
		//   channel "" is encrypted and authenticated
		//   channel "AAD" is authenticated
		ef.ChannelPut( "AAD", (const byte*)adata.data(), adata.size() );
		ef.ChannelMessageEnd( "AAD" );

		// Authenticated data *must* be pushed before
		//  Confidential/Authenticated data. Otherwise
		//  we must catch the BadState exception
		ef.ChannelPut( "", (const byte*)pdata.data(), pdata.size() );
		ef.ChannelMessageEnd( "" );

		// Pretty print
		encoded.clear();
		StringSource( cipher, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "cipher: " << encoded << endl;
	}
	catch( CryptoPP::BufferedTransformation::NoChannelSupport& e )
	{
		cerr << "Caught NoChannelSupport..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Section 3
	\*********************************/

	// Attack the first and last byte
	//if( cipher.size() > 1 )
	//{
	//  cipher[ 0 ] ^= 0x01;
	//  cipher[ cipher.size()-1 ] ^= 0x01;
	//}

	/*********************************\
		Section 4
	\*********************************/

	try
	{
		EAX< AES >::Decryption d;
		d.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		// Break the cipher text out into it's
		//  components: Encrypted Data and MAC Value
		string enc = cipher.substr( 0, cipher.length()-TAG_SIZE );
		string mac = cipher.substr( cipher.length()-TAG_SIZE );

		// Sanity checks
		assert( cipher.size() == enc.size() + mac.size() );
		assert( enc.size() == pdata.size() );
		assert( TAG_SIZE == mac.size() );

		// Not recovered - sent via clear channel
		radata = adata;

		AuthenticatedDecryptionFilter df( d, NULL,
			AuthenticatedDecryptionFilter::MAC_AT_BEGIN |
			AuthenticatedDecryptionFilter::THROW_EXCEPTION, TAG_SIZE );

		// The order of the following calls are important
		df.ChannelPut( "", (const byte*)mac.data(), mac.size() );
		df.ChannelPut( "AAD", (const byte*)adata.data(), adata.size() ); 
		df.ChannelPut( "", (const byte*)enc.data(), enc.size() );

		// If the object throws, it will most likely occur
		//  during ChannelMessageEnd()
		df.ChannelMessageEnd( "AAD" );
		df.ChannelMessageEnd( "" );

		// If the object does not throw, here's the only
		//  opportunity to check the data's integrity
		bool b = false;
		b = df.GetLastResult();
		assert( true == b );

		// Remove data from channel
		string retrieved;
		size_t n = (size_t)-1;

		// Plain text recovered from enc.data()
		df.SetRetrievalChannel( "" );
		n = (size_t)df.MaxRetrievable();
		retrieved.resize( n );

		if( n > 0 ) { df.Get( (byte*)retrieved.data(), n ); }
		rpdata = retrieved;
		assert( rpdata == pdata );

		// All is well - work with data

		// Pretty print
		encoded.clear();
		StringSource( radata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "adata (received): " << encoded << endl;

		encoded.clear();
		StringSource( rpdata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "pdata (recovered): " << encoded << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::HashVerificationFilter::HashVerificationFailed& e )
	{
		cerr << "Caught HashVerificationFailed..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Section 5
	\*********************************/

	try
	{
		EAX< AES >::Decryption d;
		d.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		// Not recovered - sent via clear channel
		radata = adata;

		AuthenticatedDecryptionFilter df( d, NULL,
			AuthenticatedDecryptionFilter::MAC_AT_END |
			AuthenticatedDecryptionFilter::THROW_EXCEPTION, TAG_SIZE );

		StringSource ss(cipher, false, new Redirector(df));

		// The order of the following calls are important
		df.ChannelPutMessageEnd( "AAD", (const byte*)adata.data(), adata.size() );
		ss.PumpAll();

		// If the object does not throw, here's the only
		//  opportunity to check the data's integrity
		bool b = false;
		b = df.GetLastResult();
		assert( true == b );

		// Remove data from channel
		string retrieved;
		size_t n = (size_t)-1;

		// Plain text recovered from enc.data()
		df.SetRetrievalChannel( "" );
		n = (size_t)df.MaxRetrievable();
		retrieved.resize( n );

		if( n > 0 ) { df.Get( (byte*)retrieved.data(), n ); }
		rpdata = retrieved;
		assert( rpdata == pdata );

		// All is well - work with data

		// Pretty print
		encoded.clear();
		StringSource( radata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "adata (received): " << encoded << endl;

		encoded.clear();
		StringSource( rpdata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "pdata (recovered from string): " << encoded << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::HashVerificationFilter::HashVerificationFailed& e )
	{
		cerr << "Caught HashVerificationFailed..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Section 6
	\*********************************/

	const string CIPHER_FILE_NAME("cipher.txt");
	ofstream cipher_txt(CIPHER_FILE_NAME);
	cipher_txt << cipher;
	cipher_txt.close();

	/*********************************\
		Section 7
	\*********************************/

	try
	{
		EAX< AES >::Decryption d;
		d.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		// Not recovered - sent via clear channel
		radata = adata;

		AuthenticatedDecryptionFilter df( d, NULL,
			AuthenticatedDecryptionFilter::MAC_AT_END |
			AuthenticatedDecryptionFilter::THROW_EXCEPTION, TAG_SIZE );

		FileSource fs(CIPHER_FILE_NAME.c_str(), false, new Redirector(df));

		// The order of the following calls are important
		df.ChannelPutMessageEnd( "AAD", (const byte*)adata.data(), adata.size() );
		fs.PumpAll();

		// If the object does not throw, here's the only
		//  opportunity to check the data's integrity
		bool b = false;
		b = df.GetLastResult();
		assert( true == b );

		// Remove data from channel
		string retrieved;
		size_t n = (size_t)-1;

		// Plain text recovered from enc.data()
		df.SetRetrievalChannel( "" );
		n = (size_t)df.MaxRetrievable();
		retrieved.resize( n );

		if( n > 0 ) { df.Get( (byte*)retrieved.data(), n ); }
		rpdata = retrieved;
		assert( rpdata == pdata );

		// All is well - work with data

		// Pretty print
		encoded.clear();
		StringSource( radata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "adata (received): " << encoded << endl;

		encoded.clear();
		StringSource( rpdata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "pdata (recovered from file): " << encoded << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::HashVerificationFilter::HashVerificationFailed& e )
	{
		cerr << "Caught HashVerificationFailed..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Section 8
	\*********************************/

	try
	{
		EAX< AES >::Decryption d;
		d.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		// Not recovered - sent via clear channel
		radata = adata;

		AuthenticatedDecryptionFilterBugFixed df( d, NULL,
			AuthenticatedDecryptionFilter::MAC_AT_END |
			AuthenticatedDecryptionFilter::THROW_EXCEPTION, TAG_SIZE );

		FileSource fs(CIPHER_FILE_NAME.c_str(), false, new Redirector(df));

		// The order of the following calls are important
		df.ChannelPutMessageEnd( "AAD", (const byte*)adata.data(), adata.size() );
		fs.PumpAll();

		// If the object does not throw, here's the only
		//  opportunity to check the data's integrity
		bool b = false;
		b = df.GetLastResult();
		assert( true == b );

		// Remove data from channel
		string retrieved;
		size_t n = (size_t)-1;

		// Plain text recovered from enc.data()
		df.SetRetrievalChannel( "" );
		n = (size_t)df.MaxRetrievable();
		retrieved.resize( n );

		if( n > 0 ) { df.Get( (byte*)retrieved.data(), n ); }
		rpdata = retrieved;
		assert( rpdata == pdata );

		// All is well - work with data

		// Pretty print
		encoded.clear();
		StringSource( radata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "adata (received): " << encoded << endl;

		encoded.clear();
		StringSource( rpdata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "pdata (recovered from file - bug fixed 1): " << encoded << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::HashVerificationFilter::HashVerificationFailed& e )
	{
		cerr << "Caught HashVerificationFailed..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Section 9
	\*********************************/

	try
	{
		EAX< AES >::Decryption d;
		d.SetKeyWithIV( key, key.size(), iv, sizeof(iv) );

		// Not recovered - sent via clear channel
		radata = adata;

		AuthenticatedDecryptionFilterBugFixed2 df( d, NULL,
			AuthenticatedDecryptionFilter::MAC_AT_END |
			AuthenticatedDecryptionFilter::THROW_EXCEPTION, TAG_SIZE );

		FileSource fs(CIPHER_FILE_NAME.c_str(), false, new Redirector(df));

		// The order of the following calls are important
		df.ChannelPutMessageEnd( "AAD", (const byte*)adata.data(), adata.size() );
		fs.PumpAll();

		// If the object does not throw, here's the only
		//  opportunity to check the data's integrity
		bool b = false;
		b = df.GetLastResult();
		assert( true == b );

		// Remove data from channel
		string retrieved;
		size_t n = (size_t)-1;

		// Plain text recovered from enc.data()
		df.SetRetrievalChannel( "" );
		n = (size_t)df.MaxRetrievable();
		retrieved.resize( n );

		if( n > 0 ) { df.Get( (byte*)retrieved.data(), n ); }
		rpdata = retrieved;
		assert( rpdata == pdata );

		// All is well - work with data

		// Pretty print
		encoded.clear();
		StringSource( radata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "adata (received): " << encoded << endl;

		encoded.clear();
		StringSource( rpdata, true, new HexEncoder( new StringSink( encoded ) ) );
		cout << "pdata (recovered from file - bug fixed 2): " << encoded << endl;
	}
	catch( CryptoPP::InvalidArgument& e )
	{
		cerr << "Caught InvalidArgument..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::AuthenticatedSymmetricCipher::BadState& e )
	{
		cerr << "Caught BadState..." << endl;
		cerr << e.what() << endl;
	}
	catch( CryptoPP::HashVerificationFilter::HashVerificationFailed& e )
	{
		cerr << "Caught HashVerificationFailed..." << endl;
		cerr << e.what() << endl;
	}

	/*********************************\
		Final Section
	\*********************************/

	return 0;
}