// Windows
#include <Windows.h>

// CRT
#include <stdio.h>
#include <stdlib.h>

// FlowSshC
#include "FlowSshC.h"

#include <iostream>



enum ExitCodes
{
	Success,		// no exit code received from remote
	UsageError,
	SessionError,	// including connection error 
	FatalError,
	ExitCodeBase = 1000,	// ExitCode = ExitCodeBase + exit code received from remote
};

volatile int g_exitCode = Success; 


void ErrorHandler(void* handlerData, unsigned int flags, wchar_t const* desc)
{
	wprintf(L"Fatal error: %s\n", desc ? desc : L"<no error description>");
	// For the sake of simplicity, no cleanup is performed here.
	// The lack of cleanup may cause debugger to report memory leaks.
	exit(FatalError);
}



struct Events
{
	HANDLE m_readyEvent; // auto reset event
	HANDLE m_exitEvent;
};


void SshVersionHandler(void* handlerData, wchar_t const* version)
{
	wprintf(L"Server version: %s\n", version);
}

bool HostKeyHandler(void* handlerData, FlowSshC_PublicKey* publicKey)
{
	bool success = false;

	BSTR md5 = FlowSshC_PublicKey_GetMd5(publicKey);
	BSTR bb = FlowSshC_PublicKey_GetBubbleBabble(publicKey);
	BSTR sha256 = FlowSshC_PublicKey_GetSha256(publicKey);

	wprintf(L"Received the following host key:\n"
			L"MD5 Fingerprint: %s\n"
			L"Bubble-Babble: %s\n"
			L"SHA-256: %s\n",
			md5, bb, sha256);

	SysFreeString(md5);
	SysFreeString(bb);
	SysFreeString(sha256);

	wprintf(L"Accept the key (yes/no)? ");

	wchar_t input[10];
	if (_getws_s(input, 9))
		if (!_wcsicmp(input, L"yes"))
			success = true;

	return success;
}

bool UserAuthHandler(void* handlerData, FlowSshC_FurtherAuth* furtherAuth, FlowSshC_PasswordChange* passwordChange)
{
	bool success = false;
	
	// Disable echo on input
	HANDLE stdIn = GetStdHandle(STD_INPUT_HANDLE);
	DWORD originalMode;
	BOOL restoreMode = FALSE;
	if (GetFileType(stdIn) == FILE_TYPE_CHAR)
		if (restoreMode = GetConsoleMode(stdIn, &originalMode))
			restoreMode = SetConsoleMode(stdIn, originalMode & ~((DWORD) ENABLE_ECHO_INPUT));

	if (furtherAuth)
	{
		if (FlowSshC_FurtherAuth_IsPasswordRemaining(furtherAuth))
		{
			wprintf(L"Password: ");	

			wchar_t password[100];	
			if (_getws_s(password, 99))
			{
				FlowSshC_FurtherAuth_SetPassword(furtherAuth, password);
				success = true;
			}

			memset(password, 0, sizeof(password));

			wprintf(L"\n");
		}
	}
	else if (passwordChange)
	{
		BSTR prompt = FlowSshC_PasswordChange_GetPrompt(passwordChange);
		// Production code should remove volatile characters from prompt first.
		wprintf(L"%s\n", prompt);
		SysFreeString(prompt);

		wprintf(L"New password: ");

		wchar_t newPassword[100];
		if (_getws_s(newPassword, 99))
		{
			wprintf(L"\nRepeat password: ");
			
			wchar_t repeatPassword[100];	
			if (_getws_s(repeatPassword, 99))
			{
				if (wcscmp(newPassword, repeatPassword))
					wprintf(L"\nPasswords do not match!");
				else
				{
					FlowSshC_PasswordChange_SetNewPassword(passwordChange, newPassword);
					success = true;
				}
			}

			memset(repeatPassword, 0, sizeof(repeatPassword));
		}

		memset(newPassword, 0, sizeof(newPassword));

		wprintf(L"\n");
	}

	// Restore echo on input
	if (restoreMode)
		SetConsoleMode(stdIn, originalMode);

	return success;
}

void BannerHandler(void* handlerData, wchar_t const* banner)
{
	// Production code should remove volatile characters from banner first.
	wprintf(L"User authentication banner:\n%s\n", banner);
}

void DisconnectHandler(void* handlerData, unsigned int reason, wchar_t const* desc)
{
	wprintf(L"%s", desc);

	if (handlerData)
	{
		if (g_exitCode == Success)
			g_exitCode = SessionError;

		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}

void ConnectProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed)
	{
		SetEvent(((Events*) handlerData)->m_readyEvent);
	}
	else if (taskState == FlowSshC_TaskState_Failed)
	{
		if (taskSpecificStep == FlowSshC_ConnectStep_ConnectToProxy)
			wprintf(L"Connecting to proxy server failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		else if (taskSpecificStep == FlowSshC_ConnectStep_ConnectToSshServer)
			wprintf(L"Connecting to SSH server failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		else if (taskSpecificStep == FlowSshC_ConnectStep_SshVersionString)
			wprintf(L"SSH version string failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		else if (taskSpecificStep == FlowSshC_ConnectStep_SshKeyExchange)
			wprintf(L"SSH key exchange failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		else if (taskSpecificStep == FlowSshC_ConnectStep_SshUserAuth)
			wprintf(L"SSH authentication failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		else
			wprintf(L"Connecting failed at unknown step: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		
		if (g_exitCode == Success)
			g_exitCode = SessionError;

		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}


void CloseHandler(void* handlerData)
{
	SetEvent(((Events*) handlerData)->m_exitEvent);
}

void ExitHandler(void* handlerData, FlowSshC_ExitStatus const* status, FlowSshC_ExitSignal const* signal)
{
	if (status)
		if (g_exitCode == Success)
			g_exitCode = ExitCodeBase + status->m_code;
}

void OpenChannelProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed)
	{
		SetEvent(((Events*) handlerData)->m_readyEvent);
	}
	else if (taskState == FlowSshC_TaskState_Failed)
	{
		wprintf(L"Opening session channel failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		
		if (g_exitCode == Success)
			g_exitCode = SessionError;

		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}

void SendPtyReqProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed)
	{
		SetEvent(((Events*) handlerData)->m_readyEvent);
	}
	else if (taskState == FlowSshC_TaskState_Failed)
	{
		wprintf(L"Pty request failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");

		if (g_exitCode == Success)
			g_exitCode = SessionError;
		
		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}

void SendExecRequestProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed)
	{
		SetEvent(((Events*) handlerData)->m_readyEvent);
	}
	else if (taskState == FlowSshC_TaskState_Failed)
	{
		wprintf(L"Exec request failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		
		if (g_exitCode == Success)
			g_exitCode = SessionError;
		
		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}

void SendShellRequestProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed)
	{
		SetEvent(((Events*) handlerData)->m_readyEvent);
	}
	else if (taskState == FlowSshC_TaskState_Failed)
	{
		wprintf(L"Shell request failed: %s\n",  auxInfo ? auxInfo : L"(no additional info)");
		
		if (g_exitCode == Success)
			g_exitCode = SessionError;
		
		SetEvent(((Events*) handlerData)->m_exitEvent);
	}
}


void DisconnectProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState == FlowSshC_TaskState_Completed || taskState == FlowSshC_TaskState_Failed)
		SetEvent((HANDLE) handlerData);
}

void CloseProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	int i = 0;
}



// Channel input thread

struct InputParams
{
	FlowSshC_ClientSessionChannel* m_channel;
	HANDLE m_thisThread;
	HANDLE m_outputThread;

	HANDLE m_sentEvent;		// auto reset event
	bool m_sentResult;
};

struct ReadParams
{
	HANDLE m_file;
	HANDLE m_thread;
	HANDLE m_readyEvent;
	HANDLE m_continueEvent;
	unsigned char m_data[4*1024];
	unsigned int m_dataBytes;
};


DWORD WINAPI ReadFileThread(LPVOID param)
{
	ReadParams* p = (ReadParams*) param;

	while (true)
	{	
		DWORD read;
		if (!ReadFile(p->m_file, p->m_data, sizeof(p->m_data), &read, 0))
			return GetLastError();

		p->m_dataBytes = read;

		SetEvent(p->m_readyEvent);
		WaitForSingleObject(p->m_continueEvent, INFINITE);
	}
	
	return 0;
};


void SendProgressHandler(void* handlerData, unsigned int taskState, unsigned int taskSpecificStep, wchar_t const* auxInfo)
{
	if (taskState != FlowSshC_TaskState_InProgress)
	{
		InputParams* p = (InputParams*) handlerData;
		p->m_sentResult = taskState == FlowSshC_TaskState_Completed;
		SetEvent(p->m_sentEvent);
	}
}


DWORD WINAPI ChannelInputThread(LPVOID param)
{
	InputParams* p = (InputParams*) param;

	ReadParams r = { 0, 0, 0, 0 };
	r.m_file = GetStdHandle(STD_INPUT_HANDLE);
	r.m_readyEvent = CreateEvent(0, false, false, 0);
	r.m_continueEvent = CreateEvent(0, false, false, 0);

	if (!r.m_readyEvent && !r.m_continueEvent)
	{
		wprintf(L"CreateEvent() failed with Windows error %u.\n", GetLastError());
		
		if (g_exitCode == Success)
			g_exitCode = FatalError;
	}
	else if (!(r.m_thread = CreateThread(0, 0, ReadFileThread, &r, 0, 0)))
	{
		wprintf(L"CreateThread() failed with Windows error %u.\n", GetLastError());
		
		if (g_exitCode == Success)
			g_exitCode = FatalError;
	}
	else
	{
		while (true)
		{
			HANDLE waitObjects[3] = { p->m_outputThread, r.m_thread, r.m_readyEvent };
			DWORD dw = WaitForMultipleObjects(3, waitObjects, false, INFINITE);
			if (dw == WAIT_OBJECT_0)
				break;
			else if (dw == WAIT_OBJECT_0 + 1)
			{
				DWORD exitCode = 0;
				GetExitCodeThread(r.m_thread, &exitCode);
				wprintf(L"ReadFile() failed with Windows error %u.\n", exitCode);
				
				if (g_exitCode == Success)
					g_exitCode = FatalError;
				break;
			}
			
			if (!r.m_dataBytes)
				break;

			FlowSshC_ClientSessionChannel_Send(p->m_channel, r.m_data, r.m_dataBytes, false, SendProgressHandler, p);
			SetEvent(r.m_continueEvent);
			WaitForSingleObject(p->m_sentEvent, INFINITE);
		}
	}

	FlowSshC_ClientSessionChannel_Send(p->m_channel, 0, 0, true, 0, 0);

	if (r.m_thread)
	{
		TerminateThread(r.m_thread, 0);
		WaitForSingleObject(r.m_thread, INFINITE);
	}

	CloseHandle(r.m_readyEvent);
	CloseHandle(r.m_continueEvent);
	CloseHandle(r.m_thread);

	return 0;
}



// Channel output thread

struct OutputParams
{
	FlowSshC_ClientSessionChannel* m_channel;
	HANDLE m_thisThread;

	HANDLE m_recvEvent;		// auto reset event
	bool m_recvStdErr;
	unsigned char* m_recvData;
	unsigned int m_recvDataBytes;	
	bool m_recvEof;
};


void ReceiveHandler(void* handlerData, bool stdErr, unsigned char const* data, unsigned int dataBytes, bool eof)
{
	OutputParams* p = (OutputParams*) handlerData;
	
	p->m_recvStdErr = stdErr;
	p->m_recvDataBytes = dataBytes;
	p->m_recvEof = eof;
	if (dataBytes) 
		memcpy(p->m_recvData, data, dataBytes);
	
	SetEvent(p->m_recvEvent);
}


DWORD WINAPI ChannelOutputThread(LPVOID param)
{
	OutputParams* p = (OutputParams*) param;
	
	unsigned char buffer[4*1024];
	p->m_recvData = buffer;
	MemoryBarrier();

	HANDLE stdErr = GetStdHandle(STD_ERROR_HANDLE);
	HANDLE stdOut = GetStdHandle(STD_OUTPUT_HANDLE);

	bool writeFileError = false;

	do
	{
		FlowSshC_ClientSessionChannel_Receive(p->m_channel, sizeof(buffer), ReceiveHandler, p);
		WaitForSingleObject(p->m_recvEvent, INFINITE);

		unsigned char* data = p->m_recvData;
		DWORD dataBytes = p->m_recvDataBytes;

		while (dataBytes && !writeFileError)
		{
			DWORD written;
			if (WriteFile(p->m_recvStdErr ? stdErr : stdOut, data, dataBytes, &written, 0))
			{
				data += written;
				dataBytes -= written;
			}
			else
			{
				wprintf(L"WriteFile() failed with Windows error %u.\n", GetLastError());
				
				if (g_exitCode == Success)
					g_exitCode = FatalError;
				writeFileError = true;
			}
		}
	}
	while (!writeFileError && p->m_recvDataBytes && !p->m_recvEof);

	return 0;
}



int wmain(int argc, wchar_t const* argv[])
{
	FlowSshC_Initialize(ErrorHandler, 0);

	// 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.
	//
	//FlowSshC_SetActCode(L"Enter Your Activation Code Here");

	FlowSshC_Client* client = FlowSshC_Client_Create();
	// CHANGE APPLICATION NAME IN PRODUCTION CODE!
	FlowSshC_Client_SetAppName(client, L"FlowSshC_Exec 1.0");

	// Process command-line parameters

	bool usageError = false;
	wchar_t const* cmd = 0;
	wchar_t const* pty = 0;
	
	for (int i = 1; i < argc && !usageError; ++i)
	{
		if (wcsncmp(argv[i], L"-user=", 6) == 0)
			FlowSshC_Client_SetUserName(client, argv[i] + 6);
		else if (wcsncmp(argv[i], L"-pw=", 4) == 0)
			FlowSshC_Client_SetPassword(client, argv[i] + 4);
		else if (wcsncmp(argv[i], L"-host=", 6) == 0)
			FlowSshC_Client_SetHost(client, argv[i] + 6);
		else if (wcsncmp(argv[i], L"-port=", 6) == 0)
		{
			unsigned int port;
			if (swscanf_s(argv[i] + 6, L"%u", &port))
				FlowSshC_Client_SetPort(client, port);
		}
		else if (wcsncmp(argv[i], L"-pty=", 5) == 0)
			pty = argv[i] + 5;
		else if (wcsncmp(argv[i], L"-cmd=", 5) == 0)
			cmd = argv[i] + 5;
		else
			usageError = true;
	}
			
	if (usageError)
	{
		wprintf(L"FlowSshC sample Exec client.\n"
				L"\n"
				L"Parameters:\n"
				L" -host=...  (default localhost)\n"
				L" -port=...  (default 22)\n"
				L" -user=...\n"
				L" -pw=...\n"
				L" -pty=...\n"
				L" -cmd=...   (open shell by default)\n"
				L"\n");
		
		if (g_exitCode == Success)
			g_exitCode = UsageError;
	}
	else
	{
		// Create events

		Events events;
		events.m_readyEvent = CreateEvent(0, false, false, 0);
		events.m_exitEvent = CreateEvent(0, true, false, 0);
		
		if (!events.m_exitEvent || !events.m_readyEvent)
		{
			wprintf(L"CreateEvent() failed with Windows error %u.\n", GetLastError());
			if (g_exitCode == Success)
				g_exitCode = FatalError;
		}
		else
		{
			// Connect the client

			FlowSshC_Client_SetVersionHandler(client, SshVersionHandler, 0);
			FlowSshC_Client_SetHostKeyHandler(client, HostKeyHandler, 0);
			FlowSshC_Client_SetUserAuthHandler(client, UserAuthHandler, 0);
			FlowSshC_Client_SetBannerHandler(client, BannerHandler, 0);
			FlowSshC_Client_SetDisconnectHandler(client, DisconnectHandler, &events);
			FlowSshC_Client_Connect(client, ConnectProgressHandler, &events);

			HANDLE readyExitHandles[2] = { events.m_readyEvent, events.m_exitEvent };
			DWORD dw = WaitForMultipleObjects(2, readyExitHandles, false, INFINITE);
			if (dw == WAIT_OBJECT_0)
			{
				// Open session channel

				FlowSshC_ClientSessionChannel* channel = FlowSshC_ClientSessionChannel_Create(client);
				FlowSshC_ClientSessionChannel_SetCloseHandler(channel, CloseHandler, &events);
				FlowSshC_ClientSessionChannel_SetExitHandler(channel, ExitHandler, 0);
				
				OutputParams outputParams = { channel, 0, 0  };
				InputParams inputParams = { channel, 0, 0, 0 };

				do // ... while (false)
				{
					FlowSshC_ClientSessionChannel_OpenRequest(channel, OpenChannelProgressHandler, &events);
					dw = WaitForMultipleObjects(2, readyExitHandles, false, INFINITE);
					if (dw != WAIT_OBJECT_0)
						break;

					// Data can arrive at any point after channel is opened.
					// Creating thread for reading channel's incoming data.
					outputParams.m_recvEvent = CreateEvent(0, false, false, 0); 
					if (!outputParams.m_recvEvent)
					{
						wprintf(L"CreateEvent() failed with Windows error %u.\n", GetLastError());
						if (g_exitCode == Success)
							g_exitCode = FatalError;
						break;
					}
					else
					{
						outputParams.m_thisThread = CreateThread(0, 0, ChannelOutputThread, &outputParams, 0, 0);
						if (!outputParams.m_thisThread)
						{
							wprintf(L"CreateThread() failed with Windows error %u.\n", GetLastError());
							if (g_exitCode == Success)
								g_exitCode = FatalError;
							break;
						}
					}

					if (pty)
					{
						FlowSshC_ClientSessionChannel_PtyRequest(channel, pty, 80, 25, SendPtyReqProgressHandler, &events);
						dw = WaitForMultipleObjects(2, readyExitHandles, false, INFINITE);
						if (dw != WAIT_OBJECT_0)
							break;
					}

					if (cmd)
						FlowSshC_ClientSessionChannel_ExecRequest(channel, cmd, SendExecRequestProgressHandler, &events);
					else
						FlowSshC_ClientSessionChannel_ShellRequest(channel, SendShellRequestProgressHandler, &events);
					
					dw = WaitForMultipleObjects(2, readyExitHandles, false, INFINITE);
					if (dw != WAIT_OBJECT_0)
						break;

					// Creating thread for writing channel's outgoing data.
					inputParams.m_sentEvent = CreateEvent(0, false, false, 0);
					if (!inputParams.m_sentEvent)
					{
						wprintf(L"CreateEvent() failed with Windows error %u.\n", GetLastError());
						if (g_exitCode == Success)
							g_exitCode = FatalError;
						break;
					}
					else
					{
						inputParams.m_outputThread = outputParams.m_thisThread;
						inputParams.m_thisThread = CreateThread(0, 0, ChannelInputThread, &inputParams, 0, 0);
						if (!inputParams.m_thisThread)
						{
							wprintf(L"CreateThread() failed with Windows error %u.\n", GetLastError());
							if (g_exitCode == Success)
								g_exitCode = FatalError;
							break;
						}
					}

					// Output thread exits if channel gets closed or if EOF is received.
					// Input thread exits if output thread has finished.
					HANDLE threadsExitHandles[2] = { inputParams.m_thisThread, outputParams.m_thisThread };
					WaitForMultipleObjects(2, threadsExitHandles, true, INFINITE);	
				}
				while (false);

				FlowSshC_ClientSessionChannel_SetCloseHandler(channel, 0, 0);
				FlowSshC_ClientSessionChannel_Close(channel, 0 , 0);
				
				// Wait output and input threads to finish gracefully.
				if (outputParams.m_thisThread)
					WaitForSingleObject(outputParams.m_thisThread, INFINITE);
				if (inputParams.m_thisThread)
					WaitForSingleObject(inputParams.m_thisThread, INFINITE);

				FlowSshC_ClientSessionChannel_Release(channel);

				if (outputParams.m_recvEvent)
					CloseHandle(outputParams.m_recvEvent);
				if (outputParams.m_thisThread)
					CloseHandle(outputParams.m_thisThread);

				if (inputParams.m_sentEvent)
					CloseHandle(inputParams.m_sentEvent);
				if (inputParams.m_thisThread)
					CloseHandle(inputParams.m_thisThread);
			}
		}

		FlowSshC_Shutdown();

		if (events.m_exitEvent)
			CloseHandle(events.m_exitEvent);
		if (events.m_readyEvent)
			CloseHandle(events.m_readyEvent);
	}

	FlowSshC_Client_Release(client);
	return g_exitCode;
}
