// Windows
#include <Windows.h>

// CRT
#include <iostream>

// FlowSshCpp
#include "FlowSshCpp.h"

// name spaces
using namespace std;
using namespace FlowSshCpp;


enum ExitCodes
{
	Success,
	UsageError,
	SessionError, // including connection error
	FatalError	
};


// CmdLineParams 

struct CmdLineParams
{
	bool		m_bOk;
	wstring		m_sUserName;
	wstring		m_sPassword;
	unsigned int m_nPort;
	wstring		m_sHost;

	bool		m_bDownload;
	wstring		m_sLocalPath;
	wstring		m_sRemotePath;

	CmdLineParams(int argc, wchar_t const* argv[]);
};

CmdLineParams::CmdLineParams(int argc, wchar_t const* argv[]) : m_bOk(false), m_nPort(0), m_bDownload(false)
{
	bool syntaxError = false;
	for (int i = 1; i < argc && !syntaxError; ++i)
	{
		if (wcsncmp(argv[i], L"-l=", 3) == 0)
			m_sUserName = argv[i] + 3;
		else if (wcsncmp(argv[i], L"-pw=", 4) == 0)
			m_sPassword = argv[i] + 4;
		else if (wcsncmp(argv[i], L"-P=", 3) == 0)
		{
			unsigned int port;
			if (swscanf_s(argv[i] + 3, L"%u", &port))
				m_nPort = port;
		}
		else 
		{
			wstring ws(argv[i]);
			size_t atpos = ws.find(L"@");
			bool bat = (atpos != wstring::npos);

			if (bat)
			{
				m_sUserName = ws.substr(0, atpos);
				ws = ws.substr(atpos + 1);
			}

			size_t colon_pos = ws.find(L":");
			bool bcol = (colon_pos != wstring::npos);
			if (bat && !bcol)
				syntaxError = true;
			else
			{
				// Note that absolute Windows paths have colon at 2nd position.
				if (bat || (bcol && colon_pos != 1)) 
				{
					m_sHost = ws.substr(0, colon_pos);
					m_sRemotePath = ws.substr(colon_pos + 1);
					if (m_sLocalPath.length() == 0)
						m_bDownload = true;
				}
				else
				{
					m_sLocalPath = ws;
					if (m_sRemotePath.length() == 0)
						m_bDownload = false;
				}
			}
		}
	}

	if (syntaxError || m_sLocalPath.length() == 0 || m_sRemotePath.length() == 0)
	{
		wcout << L"FlowSshCpp sample SFTP client." << endl;
		wcout << endl;
		wcout << L"Usage: FlowSshCpp_Sftp [Options] [User@]Host:SourceFile TargetFile" << endl;
		wcout << L"       FlowSshCpp_Sftp [Options] SourceFile [User@]Host:TargetFile" << endl;
		wcout << endl;
		wcout << L"Options:" << endl;
		wcout << L"       -l=user" << endl;
		wcout << L"       -pw=password" << endl;
		wcout << L"       -P=port" << endl;
		m_bOk = false;
	}
	else m_bOk = true;
}


// MyErrorHandler

class MyErrorHandler : public ErrorHandler
{
private:
	// This object must be created on the heap.
	~MyErrorHandler() {}

protected:
	virtual void OnExceptionInHandler(bool fatal, wchar_t const* desc)
	{
		wcout << (fatal ? L"Error [fatal]: " : L"Error: ");
		wcout << (desc ? desc : L"<no error description>") << endl;
		// For the sake of simplicity, no cleanup is performed here.
		// The lack of cleanup may cause debugger to report memory leaks.
		exit(FatalError);
	}
};


// MyClient

class MyClient : public Client
{
private:
	// This object must be created on the heap.
	~MyClient() {}

public:
	void SetClientSettings(CmdLineParams& params)
	{
		if (params.m_sUserName.size()) SetUserName(params.m_sUserName.c_str());
		if (params.m_sPassword.size()) SetPassword(params.m_sPassword.c_str());
		if (params.m_nPort)			   SetPort(params.m_nPort);
		if (params.m_sHost.size())	   SetHost(params.m_sHost.c_str());
	}

protected:
	// authenticates the server
	bool OnHostKey(RefPtr<PublicKey> publicKey)
	{
		wcout << L"Received the following host key:" << endl;
		wcout << L"  MD5 Fingerprint: " << publicKey->GetMd5() << endl;
		wcout << L"  Bubble-Babble: " << publicKey->GetBubbleBabble() << endl;
		wcout << L"  SHA-256: " << publicKey->GetSha256() << endl;
		wcout << L"Accept the key (yes/no)? ";

		std::wstring input;
		wcin >> input;

		if (input == L"yes")
			return true;

		// host key rejected
		return false;
	}
};


// MyTransferEvent - implements logging of error and success

class MyTransferEvent : public TransferEvent
{
private:
	// This object must be created on the heap.
	~MyTransferEvent() {}

protected:
	virtual void OnDone()
	{
		if (Success())
			wcout << L"File transfer completed." << endl;

		TransferEvent::OnDone();	// trigger event
	}
	virtual void OnError(TransferErr const& error);
};

void MyTransferEvent::OnError(TransferErr const& error)
{
	std::wstring errMsg;
	if (error.m_errMsg.size()) errMsg =error.m_errMsg;
	else errMsg = L"<no error message>";

	if (error.m_transferOp < 100) // codeless errors
		wcout << L"File transfer failed: " << errMsg << endl;
	else if (error.m_transferOp < 200) // SFTP errors
	{
		switch(error.m_transferOp)
		{
		case FlowSshC_TransferOp_MakeRemoteDir:		wcout << L"Create remote directory"; break;
		case FlowSshC_TransferOp_OpenRemoteFile:	wcout << L"Opening remote file"; break;
		case FlowSshC_TransferOp_ReadRemoteFile:	wcout << L"Reading remote file"; break;
		case FlowSshC_TransferOp_WriteRemoteFile:	wcout << L"Writing remote file"; break;
		default:
			wcout << L"Transfer";
		}
		wcout << L" failed with SFTP error " << error.m_errCode << L": " << errMsg << endl;
	}
	else if (error.m_transferOp < 300) // WinAPI errors
	{
		switch (error.m_transferOp)
		{
		case FlowSshC_TransferOp_MakeLocalDir:		wcout << L"Create local directory";	break;
		case FlowSshC_TransferOp_OpenLocalFile:		wcout << L"Opening local file"; break;
		case FlowSshC_TransferOp_ReadLocalFile:		wcout << L"Reading local file"; break;
		case FlowSshC_TransferOp_WriteLocalFile:	wcout << L"Writing local file"; break;
		default:
			wcout << L"Transfer";
		}
		wcout << L" failed with Windows error " << error.m_errCode << L": " << errMsg << endl;
	}
	else // FlowSshC_ResumeErrCode
		wcout << L"Resuming file failed: " << errMsg << endl;
}


// wmain 

int wmain(int argc, wchar_t const* argv[])
{	
	CmdLineParams params(argc, argv);
	if (!params.m_bOk)
		return UsageError;
	
	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");

		// create client

		RefPtr<MyClient> client(new MyClient);
		// CHANGE APPLICATION NAME IN PRODUCTION CODE!
		client->SetAppName(L"FlowSshCpp_Sftp 1.0");
		client->SetClientSettings(params);

		// connect client

		RefPtr<ProgressEvent> progress(new ProgressEvent);
		client->Connect(progress);
		progress->WaitDone();

		if (!progress->Success())
		{	// Alternatively we could derive a class from ProgressEvent,  
			// override ProgressEvent::OnError, and do the logging there.
			std::wstring auxInfo;
			if (progress->GetAuxInfo().size()) auxInfo = progress->GetAuxInfo();
			else auxInfo = L"(no additional info)";

			switch(progress->GetTaskSpecificStep())
			{
			case FlowSshC_ConnectStep_ConnectToProxy:		wcout << L"Connecting to proxy server failed: " << auxInfo << endl; break;
			case FlowSshC_ConnectStep_ConnectToSshServer:	wcout << L"Connecting to SSH server failed: " << auxInfo << endl; break;
			case FlowSshC_ConnectStep_SshVersionString:		wcout << L"SSH version string failed: " << auxInfo << endl;	break;
			case FlowSshC_ConnectStep_SshKeyExchange:		wcout << L"SSH key exchange failed: " << auxInfo << endl; break;
			case FlowSshC_ConnectStep_SshUserAuth:			wcout << L"SSH authentication failed: " << auxInfo << endl;	break;
			default:										wcout << L"Connecting failed at unknown step: " << auxInfo << endl; break;
			}
			return SessionError;
		}

		// open sftp channel

		RefPtr<ClientSftpChannel> sftpChannel(new ClientSftpChannel(client));
		sftpChannel->Open(progress);
		progress->WaitDone();

		if (!progress->Success())
		{	// Alternatively we could derive a class from ProgressEvent,  
			// override ProgressEvent::OnError, and do the logging there.
			std::wstring auxInfo;
			if (progress->GetAuxInfo().size()) auxInfo = progress->GetAuxInfo();
			else auxInfo = L"(no additional info)";

			switch(progress->GetTaskSpecificStep())
			{
			case FlowSshC_ClientSftpChannelOpenStep_OpenRequest: wcout << L"Opening SFTP channel failed: " << endl << auxInfo << endl;	break;
			case FlowSshC_ClientSftpChannelOpenStep_SftpRequest: wcout << L"Requesting SFTP subsystem failed: " << endl << auxInfo << endl; break;
			case FlowSshC_ClientSftpChannelOpenStep_InitPacket:  wcout << L"Initializing SFTP protocol failed: " << endl << auxInfo << endl; break;
			default:											 wcout << L"Opening SFTP channel failed at unknown step: " << endl << auxInfo << endl; break;
			}
			return SessionError;
		}

		// transfer file

		RefPtr<MyTransferEvent> transfer(new MyTransferEvent);
		if (params.m_bDownload)
		{
			wcout << L"Downloading " << params.m_sRemotePath << L" to " << params.m_sLocalPath << endl;
			sftpChannel->Download(params.m_sRemotePath.c_str(), params.m_sLocalPath.c_str(), FlowSshC_TransferFlags_Binary, transfer);
		}
		else
		{
			wcout << L"Uploading " << params.m_sLocalPath << L" to " << params.m_sRemotePath << endl;
			sftpChannel->Upload(params.m_sLocalPath.c_str(), params.m_sRemotePath.c_str(), FlowSshC_TransferFlags_Binary, transfer);
		}
		transfer->WaitDone();

		// In contrast to the ProgressEvent used above we derive our own MyTransferEvent.
		// Success is logged in its MyTransferEvent::OnDone override.
		// Errors are logged in its MyTransferEvent::OnError override.
		if (!transfer->Success())
			return SessionError;
	}
	catch (Exception const& e)
	{
		wcout << e.What() << endl;
		return FatalError;
	}
	return Success;
}