// Windows
#include <Windows.h>

// CRT
#include <wchar.h>

// STL
#include <iostream>
#include <sstream>
#include <fstream>
#include <memory>

// FlowSshCpp
#include "FlowSshCpp.h"


using namespace std;
using namespace FlowSshCpp;


class MyErrorHandler : public ErrorHandler
{
private:
	~MyErrorHandler() {}

protected:
	virtual void OnExceptionInHandler(bool fatal, wchar_t const* desc) { exit(3); }
};


int wmain(int argc, wchar_t const* argv[])
{
	if (argc != 1)
	{
		wcout << L"FlowSshCpp sample Key manager." << endl;
		return 1;
	}

	try
	{
		// Initialize FlowSsh and register an ErrorHandler for uncaught exceptions in user-defined handlers.
		// Example: If there is an uncaught exception in MyClient::OnHostKey, 
		//          then this is reported in MyErrorHandler::OnExceptionInHandler.
		Initializer init(new MyErrorHandler());

		// For use in deployed applications where Bitvise SSH Client might not be installed, 
		// provide an activation code using SetActCode to ensure that FlowSsh does not  
		// display an evaluation dialog. On computers with Bitvise SSH Client, use of 
		// FlowSsh is permitted under the same terms as the Client; in this case, there 
		// will be no evaluation dialog, and SetActCode does not have to be called.
		//
		//init.SetActCode(L"Enter Your Activation Code Here");

		RefPtr<Keypair> keypair[100];
		RefPtr<PublicKey> publicKey[100];

		while (true)
		{
			wcout << L"> ";

			wstring line;
			getline(wcin, line, L'\n');
			for (wstring::iterator it = line.begin(); it != line.end(); ++it)
				*it = towlower(*it);
			
			wistringstream wsin(line);
			
			wstring cmd, type;
			unsigned int idx;
			wsin >> cmd >> type >> idx;

			if (cmd == L"quit") cmd = L"exit";
			else if (cmd == L"store") cmd = L"save";
			else if (cmd == L"about") cmd = L"info";
			
			if (type == L"publickey") type = L"public";
			else if (type == L"keypair") type = L"private";
			else if (type == L"privatekey") type = L"private";

			bool help = false;

			if (cmd == L"exit")
				break;
			else if (wsin.fail() || (type != L"private" && type != L"public") || idx >= 100)
				help = true;
			else if (cmd == L"create")
			{
				if (type == L"private")
				{
					wstring alg;
					unsigned int bits;
					wsin >> alg >> bits;

					if (alg == L"rsa") alg = L"ssh-rsa"; 
					else if (alg == L"dsa") alg = L"ssh-dss";
					else if (alg == L"dss") alg = L"ssh-dss";
					else if (alg == L"ecdsa/secp256k1") alg = L"ecdsa-sha2-1.3.132.0.10";
					else if (alg == L"ecdsa/nistp256") alg = L"ecdsa-sha2-nistp256";
					else if (alg == L"ecdsa/nistp384") alg = L"ecdsa-sha2-nistp384";
					else if (alg == L"ecdsa/nistp521") alg = L"ecdsa-sha2-nistp521";

					if (wsin.fail())
						help = true;
					else if (alg != L"ecdsa-sha2-1.3.132.0.10" && alg != L"ecdsa-sha2-nistp256" &&
							 alg != L"ecdsa-sha2-nistp384" && alg != L"ecdsa-sha2-nistp521")
					{
						if (alg != L"ssh-rsa" && alg != L"ssh-dss")
							help = true;
						else if (bits < 128 || bits > 10*1024)
							help = true;
					}

					if (!help)
						keypair[idx] = new Keypair(alg.c_str(), bits);
				}
				else
				{
					wstring from, type2;
					unsigned int idx2;
					wsin >> type2 >> idx2;
					
					if (type2 == L"keypair") type2 = L"private";
					else if (type2 == L"privatekey") type2 = L"private";
							
					if (wsin.fail() || type2 != L"private" || idx2 >= 100)
						help = true;
					else if (!keypair[idx2].Get())
						wcout << L"No private key at index " << idx2 << L"." << endl;
					else
						publicKey[idx] = new PublicKey(keypair[idx2]);
				}
			}
			else if (cmd == L"load")
			{
				wstring path;
				wsin >> path;

				if (wsin.fail() || path.empty())
					help = true;
				else
				{
					basic_ifstream<unsigned char> fin;
					fin.open(path.c_str(), ios::in | ios::binary);
					if (!fin.is_open())
						wcout << "Failed to open file \"" <<  path << L"\"." << endl;
					else
					{
						Data data(10*1024);
						fin.read(data.GetWritablePtr(), data.GetSize());
						data.Resize((unsigned int) fin.gcount());
						fin.close();

						try
						{
							if (type == L"public")
								publicKey[idx] = new PublicKey(data);
							else
							{
								wstring passphrase;
								wsin >> passphrase;
								if (wsin.fail())
									passphrase.clear();

								keypair[idx] = new Keypair(data, passphrase.c_str());
							}
						}
						catch (PublicKeyLoadException const&)
						{
							wcout << "Failed to load public key file \"" <<  path << L"\"." << endl;
						}
						catch (KeypairLoadException const&)
						{
							wcout << "Failed to load private key file \"" <<  path << L"\"." << endl;
						}
					}
				}
			}
			else if (cmd == L"save")
			{
				wstring path, format;
				wsin >> path;
				
				Data data;

				if (wsin.fail() || path.empty())
					help = true;
				else if (type == L"public")
				{
					wsin >> format;
					if (wsin.fail())
						format.clear();

					if (!publicKey[idx].Get())
						wcout << L"No public key at index " << idx << L"." << endl;
					if (format.empty() || format == L"ssh2")
						data = publicKey[idx]->GetSsh2Data();
					else if (format == L"openssh")
						data = publicKey[idx]->GetOpenSshData();
					else
						help = true;
				} 
				else
				{
					wsin >> format;
					if (wsin.fail())
						format.clear();

					if (!keypair[idx].Get())
						wcout << L"No private key at index " << idx << L"." << endl;
					if (format.empty() || format == L"bitvise")
						data = keypair[idx]->GetBitviseData();
					else if (format == L"openssh")
					{
						try 
						{
							data = keypair[idx]->GetOpenSshData(); 
						}
						catch (KeypairExportException const&)
						{
							wcout << "Unable to export private key to OpenSSH format." << endl;
						}
					}
					else if (format == L"putty")
					{
						try
						{
							data = keypair[idx]->GetPuttyData();
						}
						catch (KeypairExportException const&)
						{
							wcout << "Unable to export private key to PuTTY format." << endl;
						}
					}
					else
						help = true;
				}

				if (data.GetSize())
				{
					basic_ofstream<unsigned char> fout;
					fout.open(path.c_str(), ios::out | ios::binary);
					if (!fout.is_open())
						wcout << "Failed to open file \"" <<  path << L"\"." << endl;
					else
						fout.write(data.GetPtr(), data.GetSize());
				}
			}
			else if (cmd == L"info")
			{
				RefPtr<PublicKey> info;
				if (type == L"public")
				{
					if (!publicKey[idx].Get())
						wcout << L"No public key at index " << idx << L"." << endl;
					else
						info = publicKey[idx];
				}
				else 
				{
					if (!keypair[idx].Get())
						wcout << L"No private key at index " << idx << L"." << endl;
					else
						info = new PublicKey(keypair[idx]);
				}

				if (info.Get())
					wcout << L"Algorithm: " << info->GetAlg() << endl 
			  			  << L"Bits: " << info->GetBitCount() << endl
						  << L"MD5: " << info->GetMd5() << endl;
			}
			else if (cmd == L"passphrase")
			{
				if (type == L"public")
					help = true;
				else if (!keypair[idx].Get())
					wcout << L"No private key at index " << idx << L"." << endl;
				else
				{	
					wstring passphrase;
					wsin >> passphrase;
					if (wsin.fail())
						passphrase.clear();
					
					keypair[idx]->SetPassphrase(passphrase.c_str());
				}
			}

			if (help)
			{
				wcout << L"Available commands: " << endl
			  		  << L"create private <index> <algorithm> <bits>" << endl
					  << L"create public <index1> private <index2>" << endl
					  << L"load private <index> <file> [<passphrase>]" << endl
					  << L"load public <index> <file>" << endl
					  << L"save private <index> <file> [<private_format>]" << endl
					  << L"save public <index> <file> [<public_format>]" << endl
					  << L"info public <index>" << endl
					  << L"info private <index>" << endl
					  << L"passphrase private <index> [<new_passphrase>]" << endl
					  << endl
					  << L"Valid index range: 0-99 (inclusive)" << endl
					  << L"Valid algorithms: RSA, DSA, ECDSA/secp256k1," << endl
					  << L"  ECDSA/nistp521, ECDSA/nistp384, ECDSA/nistp256" << endl
					  << L"Valid bits range: 128-10240 (inclusive)" << endl
					  << L"Valid private formats: Bitvise, OpenSSH, PuTTY" << endl
					  << L"Valid public formats: SSH2, OpenSSH" << endl
					  << endl;
			}
		}
	}
	catch (exception const& e)
	{
		wcout << L"Exception caught: " << e.what() << endl;
		return 2;
	}

	return 0;
}
