// test.cpp - written and placed in the public domain by Wei Dai

#include "pch.h"
#include "md5.h"
#include "sha.h"
#include "ripemd.h"
#include "files.h"
#include "rng.h"
#include "hex.h"
#include "gzip.h"
#include "default.h"
#include "rsa.h"
#include "randpool.h"
#include "ida.h"
#include "base64.h"
#include "socketft.h"
#include "dsa.h"
#include "rsa.h"
#include "osrng.h"
#include "wait.h"
#include "fips140.h"

#include "validate.h"
#include "bench.h"

#include <iostream>
#include <time.h>

#if defined(_WIN32) || defined(__CYGWIN__)
#include <windows.h>
#endif

#if (_MSC_VER >= 1000)
#include <crtdbg.h>		// for the debug heap
#endif

#if defined(__MWERKS__) && defined(macintosh)
#include <console.h>
#endif

USING_NAMESPACE(CryptoPP)
USING_NAMESPACE(std)

const int MAX_PHRASE_LENGTH=250;

void GenerateRSAKey(unsigned int keyLength, const char *privFilename, const char *pubFilename, const char *seed);
string RSAEncryptString(const char *pubFilename, const char *seed, const char *message);
string RSADecryptString(const char *privFilename, const char *ciphertext);
void RSASignFile(const char *privFilename, const char *messageFilename, const char *signatureFilename);
bool RSAVerifyFile(const char *pubFilename, const char *messageFilename, const char *signatureFilename);

void DigestFile(const char *file);

string EncryptString(const char *plaintext, const char *passPhrase);
string DecryptString(const char *ciphertext, const char *passPhrase);

void EncryptFile(const char *in, const char *out, const char *passPhrase);
void DecryptFile(const char *in, const char *out, const char *passPhrase);

void SecretShareFile(int threshold, int nShares, const char *filename, const char *seed);
void SecretRecoverFile(int threshold, const char *outFilename, char *const *inFilenames);

void InformationDisperseFile(int threshold, int nShares, const char *filename);
void InformationRecoverFile(int threshold, const char *outFilename, char *const *inFilenames);

void GzipFile(const char *in, const char *out, int deflate_level);
void GunzipFile(const char *in, const char *out);

void Base64Encode(const char *in, const char *out);
void Base64Decode(const char *in, const char *out);
void HexEncode(const char *in, const char *out);
void HexDecode(const char *in, const char *out);

void ForwardTcpPort(const char *sourcePort, const char *destinationHost, const char *destinationPort);

void FIPS140_SampleApplication(const char *moduleFilename, const char *edcFilename);
void FIPS140_GenerateRandomFiles();

bool Validate(int, bool, const char *);

#ifdef __BCPLUSPLUS__
int cmain(int argc, char *argv[])
#elif defined(_MSC_VER)
int __cdecl main(int argc, char *argv[])
#else
int main(int argc, char *argv[])
#endif
{
#ifdef _CRTDBG_LEAK_CHECK_DF
	// Turn on leak-checking
	int tempflag = _CrtSetDbgFlag( _CRTDBG_REPORT_FLAG );
	tempflag |= _CRTDBG_LEAK_CHECK_DF;
	_CrtSetDbgFlag( tempflag );
#endif

#if defined(__MWERKS__) && defined(macintosh)
	argc = ccommand(&argv);
#endif

	try
	{
		std::string command, executableName, edcFilename;

		if (argc < 2)
			command = 'h';
		else
			command = argv[1];

		if (FIPS_140_2_ComplianceEnabled())
		{
			edcFilename = "edc.dat";

#if defined(_WIN32) || defined(__CYGWIN__)
			TCHAR filename[MAX_PATH];
			GetModuleFileName(GetModuleHandle(NULL), filename, sizeof(filename));
			executableName = filename;
			std::string::size_type pos = executableName.rfind('\\');
			if (pos != std::string::npos)
				edcFilename = executableName.substr(0, pos+1) + edcFilename;
#else
			executableName = argv[0];
#endif

			if (command.substr(0, 4) != "fips")
			{
				byte expectedModuleDigest[SHA1::DIGESTSIZE];
				FileSource(edcFilename.c_str(), true, new HexDecoder(new ArraySink(expectedModuleDigest, sizeof(expectedModuleDigest))));

				DoPowerUpSelfTest(executableName.c_str(), expectedModuleDigest);
			}
		}

		switch (command[0])
		{
		case 'g':
		  {
			char seed[1024], privFilename[128], pubFilename[128];
			unsigned int keyLength;

			cout << "Key length in bits: ";
			cin >> keyLength;

			cout << "\nSave private key to file: ";
			cin >> privFilename;

			cout << "\nSave public key to file: ";
			cin >> pubFilename;

			cout << "\nRandom Seed: ";
			ws(cin);
			cin.getline(seed, 1024);

			GenerateRSAKey(keyLength, privFilename, pubFilename, seed);
			return 0;
		  }
		case 'r':
		  {
			switch (argv[1][1])
			{
			case 's':
				RSASignFile(argv[2], argv[3], argv[4]);
				return 0;
			case 'v':
			  {
				bool verified = RSAVerifyFile(argv[2], argv[3], argv[4]);
				cout << (verified ? "valid signature" : "invalid signature") << endl;
				return 0;
			  }
			default:
			  {
				char privFilename[128], pubFilename[128];
				char seed[1024], message[1024];

				cout << "Private key file: ";
				cin >> privFilename;

				cout << "\nPublic key file: ";
				cin >> pubFilename;

				cout << "\nRandom Seed: ";
				ws(cin);
				cin.getline(seed, 1024);

				cout << "\nMessage: ";
				cin.getline(message, 1024);

				string ciphertext = RSAEncryptString(pubFilename, seed, message);
				cout << "\nCiphertext: " << ciphertext << endl;

				string decrypted = RSADecryptString(privFilename, ciphertext.c_str());
				cout << "\nDecrypted: " << decrypted << endl;

				return 0;
			  }
			}
		  }
		case 'm':
			DigestFile(argv[2]);
			return 0;
		case 't':
		  {
			// VC60 workaround: use char array instead of std::string to workaround MSVC's getline bug
			char passPhrase[MAX_PHRASE_LENGTH], plaintext[1024];

			cout << "Passphrase: ";
			cin.getline(passPhrase, MAX_PHRASE_LENGTH);

			cout << "\nPlaintext: ";
			cin.getline(plaintext, 1024);

			string ciphertext = EncryptString(plaintext, passPhrase);
			cout << "\nCiphertext: " << ciphertext << endl;

			string decrypted = DecryptString(ciphertext.c_str(), passPhrase);
			cout << "\nDecrypted: " << decrypted << endl;

			return 0;
		  }
		case 'e':
		case 'd':
			if (command == "e64")
				Base64Encode(argv[2], argv[3]);
			else if (command == "d64")
				Base64Decode(argv[2], argv[3]);
			else if (command == "e16")
				HexEncode(argv[2], argv[3]);
			else if (command == "d16")
				HexDecode(argv[2], argv[3]);
			else
			{
				char passPhrase[MAX_PHRASE_LENGTH];
				cout << "Passphrase: ";
				cin.getline(passPhrase, MAX_PHRASE_LENGTH);
				if (command == "e")
					EncryptFile(argv[2], argv[3], passPhrase);
				else
					DecryptFile(argv[2], argv[3], passPhrase);
			}
			return 0;
		case 's':
			if (argv[1][1] == 's')
			{
				char seed[1024];
				cout << "\nRandom Seed: ";
				ws(cin);
				cin.getline(seed, 1024);
				SecretShareFile(atoi(argv[2]), atoi(argv[3]), argv[4], seed);
			}
			else
				SecretRecoverFile(argc-3, argv[2], argv+3);
			return 0;
		case 'i':
			if (argv[1][1] == 'd')
				InformationDisperseFile(atoi(argv[2]), atoi(argv[3]), argv[4]);
			else
				InformationRecoverFile(argc-3, argv[2], argv+3);
			return 0;
		case 'v':
			return !Validate(argc>2 ? atoi(argv[2]) : 0, argv[1][1] == 'v', argc>3 ? argv[3] : NULL);
		case 'b':
			if (argc<3)
				BenchMarkAll();
			else
				BenchMarkAll((float)atof(argv[2]));
			return 0;
		case 'z':
			GzipFile(argv[3], argv[4], argv[2][0]-'0');
			return 0;
		case 'u':
			GunzipFile(argv[2], argv[3]);
			return 0;
		case 'f':
			if (command == "fips")
				FIPS140_SampleApplication(executableName.c_str(), edcFilename.c_str());
			else if (command == "fips-rand")
				FIPS140_GenerateRandomFiles();
			else if (command == "ft")
				ForwardTcpPort(argv[2], argv[3], argv[4]);
			return 0;
		default:
			FileSource usage("usage.dat", true, new FileSink(cout));
			return 1;
		}
	}
	catch(CryptoPP::Exception &e)
	{
		cout << "\nCryptoPP::Exception caught: " << e.what() << endl;
		return -1;
	}
	catch(std::exception &e)
	{
		cout << "\nstd::exception caught: " << e.what() << endl;
		return -2;
	}
}

void FIPS140_SampleApplication(const char *moduleFilename, const char *edcFilename)
{
	if (!FIPS_140_2_ComplianceEnabled())
	{
		cerr << "FIPS-140-2 compliance was turned off at compile time.\n";
		abort();
	}

	// try to use a crypto algorithm before doing a self test
	try
	{
		// trying to use a crypto algorithm before power-up self test will result in an exception
		DES::Encryption des;

		// should not be here
		cerr << "Use of DES before power-up test failed to cause an exception.\n";
		abort();
	}
	catch (SelfTestFailure &e)
	{
		cout << "0. Caught expected exception. Exception message follows: ";
		cout << e.what() << endl;
	}

	// simulate a power-up self test error
	SimulatePowerUpSelfTestFailure();
	try
	{
		// trying to use a crypto algorithm after power-up self test error will result in an exception
		DES::Encryption des;

		// should not be here
		cerr << "Use of DES failed to cause an exception after power-up self test error.\n";
		abort();
	}
	catch (SelfTestFailure &e)
	{
		cout << "1. Caught expected exception when simulating self test failure. Exception message follows: ";
		cout << e.what() << endl;
	}

	// clear the self test error state and do power-up self test
	byte expectedModuleDigest[SHA1::DIGESTSIZE];
	FileSource(edcFilename, true, new HexDecoder(new ArraySink(expectedModuleDigest, sizeof(expectedModuleDigest))));

	DoPowerUpSelfTest(moduleFilename, expectedModuleDigest);
	if (GetPowerUpSelfTestStatus() != POWER_UP_SELF_TEST_PASSED)
	{
		cerr << "Power-up self test failed.\n";
		abort();
	}
	cout << "2. Power-up self test passed.\n";

	// encrypt and decrypt
	const byte key[] = {0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef};
	const byte iv[] = {0x12,0x34,0x56,0x78,0x90,0xab,0xcd,0xef};
	const byte plaintext[] = {	// "Now is the time for all " without tailing 0
		0x4e,0x6f,0x77,0x20,0x69,0x73,0x20,0x74,
		0x68,0x65,0x20,0x74,0x69,0x6d,0x65,0x20,
		0x66,0x6f,0x72,0x20,0x61,0x6c,0x6c,0x20};
	byte ciphertext[24];
	byte decrypted[24];

	CFB_Mode<DES>::Encryption encryption_DES_CBC;
	encryption_DES_CBC.SetKeyWithIV(key, 8, iv);
	encryption_DES_CBC.ProcessString(ciphertext, plaintext, 24);

	CFB_Mode<DES>::Decryption decryption_DES_CBC;
	decryption_DES_CBC.SetKeyWithIV(key, 8, iv);
	decryption_DES_CBC.ProcessString(decrypted, ciphertext, 24);

	if (memcmp(plaintext, decrypted, 24) != 0)
	{
		cerr << "DES-CBC Encryption/decryption failed.\n";
		abort();
	}
	cout << "3. DES-CBC Encryption/decryption succeeded.\n";

	// hash
	const byte message[] = {'a', 'b', 'c'};
	const byte expectedDigest[] = {0xA9,0x99,0x3E,0x36,0x47,0x06,0x81,0x6A,0xBA,0x3E,0x25,0x71,0x78,0x50,0xC2,0x6C,0x9C,0xD0,0xD8,0x9D};
	byte digest[20];
	
	SHA1 sha;
	sha.Update(message, 3);
	sha.Final(digest);

	if (memcmp(digest, expectedDigest, 20) != 0)
	{
		cerr << "SHA-1 hash failed.\n";
		abort();
	}
	cout << "4. SHA-1 hash succeeded.\n";

	// create auto-seeded X9.17 RNG object, if available
#ifdef OS_RNG_AVAILABLE
	AutoSeededX917RNG<DES_EDE3> rng;
#else
	// this is used to allow this function to compile on platforms that don't have auto-seeded RNGs
	RandomNumberGenerator &rng(NullRNG());
#endif

	// generate DSA key
	DSA::PrivateKey dsaPrivateKey;
	dsaPrivateKey.GenerateRandomWithKeySize(rng, 1024);
	DSA::PublicKey dsaPublicKey;
	dsaPublicKey.AssignFrom(dsaPrivateKey);
	if (!dsaPrivateKey.Validate(rng, 3) || !dsaPublicKey.Validate(rng, 3))
	{
		cerr << "DSA key generation failed.\n";
		abort();
	}
	cout << "5. DSA key generation succeeded.\n";

	// encode DSA key
	std::string encodedDsaPublicKey, encodedDsaPrivateKey;
	dsaPublicKey.DEREncode(StringSink(encodedDsaPublicKey).Ref());
	dsaPrivateKey.DEREncode(StringSink(encodedDsaPrivateKey).Ref());

	// decode DSA key
	DSA::PrivateKey decodedDsaPrivateKey;
	decodedDsaPrivateKey.BERDecode(StringStore(encodedDsaPrivateKey).Ref());
	DSA::PublicKey decodedDsaPublicKey;
	decodedDsaPublicKey.BERDecode(StringStore(encodedDsaPublicKey).Ref());

	if (!decodedDsaPrivateKey.Validate(rng, 3) || !decodedDsaPublicKey.Validate(rng, 3))
	{
		cerr << "DSA key encode/decode failed.\n";
		abort();
	}
	cout << "6. DSA key encode/decode succeeded.\n";

	// sign and verify
	byte signature[40];
	DSA::Signer signer(dsaPrivateKey);
	assert(signer.SignatureLength() == 40);
	signer.SignMessage(rng, message, 3, signature);

	DSA::Verifier verifier(dsaPublicKey);
	if (!verifier.VerifyMessage(message, 3, signature))
	{
		cerr << "DSA signature and verification failed.\n";
		abort();
	}
	cout << "7. DSA signature and verification succeeded.\n";


	// try to verify an invalid signature
	signature[0] ^= 1;
	if (verifier.VerifyMessage(message, 3, signature))
	{
		cerr << "DSA signature verification failed to detect bad signature.\n";
		abort();
	}
	cout << "8. DSA signature verification successfully detected bad signature.\n";

	// try to use an invalid key length
	try
	{
		encryption_DES_CBC.SetKey(key, 5);

		// should not be here
		cerr << "DES implementation did not detect use of invalid key length.\n";