/*
 *  JoeyDev
 *  Copyright (C) 2018-2023 Scott Duensing <scott@kangaroopunch.com>
 *
 *  This software is provided 'as-is', without any express or implied
 *  warranty.  In no event will the authors be held liable for any damages
 *  arising from the use of this software.
 *
 *  Permission is granted to anyone to use this software for any purpose,
 *  including commercial applications, and to alter it and redistribute it
 *  freely, subject to the following restrictions:
 *
 *  1. The origin of this software must not be misrepresented; you must not
 *     claim that you wrote the original software. If you use this software
 *     in a product, an acknowledgment in the product documentation would be
 *     appreciated but is not required.
 *  2. Altered source versions must be plainly marked as such, and must not be
 *     misrepresented as being the original software.
 *  3. This notice may not be removed or altered from any source distribution.
*/


#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>

#include "ssh.h"
#include "utils.h"
#include "array.h"


#ifdef _WIN32
#define socketclose  closesocket
closesocket(data->sock);
#else
#define socketclose  close
#endif


static SSHT **_activeSSHs = NULL;
static char  *_output     = NULL;
static int    _outputLen  = 0;
static char  *_buffer     = NULL;
static int    _bufferLen  = 0;


gboolean   sshUpdate(gpointer userData);  // Not static
static int sshWaitSocket(int socket_fd, LIBSSH2_SESSION *session);


SSHT *sshConnect(char *hostname, uint16_t port, char *user, char *password) {
	int                 rc;
	struct hostent     *hostent;
	struct sockaddr_in  sin;
	in_addr_t           hostaddr;
	SSHT               *data = NULL;

	hostent = gethostbyname(hostname);
	if (hostent == NULL) {
		debug("gethostbyname() failed.\n");
		return NULL;
	}

	hostaddr = inet_addr(inet_ntoa(*(struct in_addr*)*(hostent->h_addr_list)));
	if (hostaddr == (in_addr_t)-1) {
		debug("inet_addr() failed.\n");
		return NULL;
	}

	data = NEW(SSHT);

	if (data) {
		data->sock = socket(AF_INET, SOCK_STREAM, 0);

		sin.sin_family      = AF_INET;
		sin.sin_port        = htons(port);
		sin.sin_addr.s_addr = hostaddr;
		if (connect(data->sock, (struct sockaddr*)(&sin), sizeof(struct sockaddr_in)) != 0) {
			debug("connect() failed:  %d\n", errno);
			DEL(data);
			return NULL;
		}

		data->session = libssh2_session_init();
		if (!data->session) {
			socketclose(data->sock);
			DEL(data);
			return NULL;
		}

		libssh2_session_set_blocking(data->session, 0);

		while ((rc = libssh2_session_handshake(data->session, data->sock)) == LIBSSH2_ERROR_EAGAIN)
			;
		if (rc) {
			debug("Failure establishing SSH session:  %d\n", rc);
			libssh2_session_free(data->session);
			socketclose(data->sock);
			DEL(data);
			return NULL;
		}

		while ((rc = libssh2_userauth_password(data->session, user, password)) == LIBSSH2_ERROR_EAGAIN)
			;
		if (rc) {
			debug("Failure authenticating SSH session:  %d\n", rc);
			while (libssh2_session_disconnect(data->session, "Error") == LIBSSH2_ERROR_EAGAIN)
				;
			libssh2_session_free(data->session);
			socketclose(data->sock);
			DEL(data);
			return NULL;
		}
	}

	return data;
}


void sshDisconnect(SSHT **sshData) {
	SSHT *data = *sshData;

	libssh2_session_disconnect(data->session, "Normal Shutdown.");
	libssh2_session_free(data->session);
	socketclose(data->sock);

	DEL(data);
}


int sshExecute(SSHT *sshData, char *command, char **output) {
	int    rc;
	int    exitcode   = 127;
	char  *exitsignal = (char *)"none";
	int    outputLen  = 0;

	utilEnsureBufferSize((unsigned char **)output, &outputLen, 1024);
	*output[0] = 0;

	while ((sshData->channel = libssh2_channel_open_session(sshData->session)) == NULL && libssh2_session_last_error(sshData->session, NULL, NULL, 0) == LIBSSH2_ERROR_EAGAIN) {
		sshWaitSocket(sshData->sock, sshData->session);
	}
	if (sshData->channel == NULL) {
		return exitcode;
	}
	while ((rc = libssh2_channel_exec(sshData->channel, command)) == LIBSSH2_ERROR_EAGAIN) {
		sshWaitSocket(sshData->sock, sshData->session);
	}
	if (rc != 0) {
		return exitcode;
	}

	for (;;) {
		do {
			rc = libssh2_channel_read(sshData->channel, (char *)_buffer, sizeof(_buffer));
			if (rc > 0) {
				utilEnsureBufferSize((unsigned char **)output, &outputLen, (int)strlen((const char *)*output) + rc + 1);
				strncat(*output, (char *)_buffer, rc);
			}
		} while (rc > 0);

		if (rc == LIBSSH2_ERROR_EAGAIN) {
			sshWaitSocket(sshData->sock, sshData->session);
		} else {
			break;
		}
	}

	while ((rc = libssh2_channel_close(sshData->channel)) == LIBSSH2_ERROR_EAGAIN) {
		sshWaitSocket(sshData->sock, sshData->session);
	}
	if (rc == 0) {
		exitcode = libssh2_channel_get_exit_status(sshData->channel);
		libssh2_channel_get_exit_signal(sshData->channel, &exitsignal, NULL, NULL, NULL, NULL, NULL);
	}

	libssh2_channel_free(sshData->channel);

	return exitcode;
}


SSHT *sshExecuteVerbose(SSHT *sshData, char *title, char *command, SSHCallback callback) {
	int rc;

	sshData->title    = strdup(title);
	sshData->callback = callback;
	sshData->finished = FALSE;

	while ((sshData->channel = libssh2_channel_open_session(sshData->session)) == NULL && libssh2_session_last_error(sshData->session, NULL, NULL, 0) == LIBSSH2_ERROR_EAGAIN) {
		sshWaitSocket(sshData->sock, sshData->session);
	}
	if (sshData->channel == NULL) {
		return NULL;
	}
	while ((rc = libssh2_channel_exec(sshData->channel, command)) == LIBSSH2_ERROR_EAGAIN) {
		sshWaitSocket(sshData->sock, sshData->session);
	}
	if (rc != 0) {
		return NULL;
	}

	//***TODO*** SSH Progress Dialog Here
	//sshData->status = new SSHStatus(sshData);
	arrpush(_activeSSHs, sshData);

	return sshData;
}


void sshShutdown(void) {
	//***TODO*** Any active sessions?

	g_idle_remove_by_data(sshUpdate);

	DEL(_output);
	DEL(_buffer);

	libssh2_exit();
}


void sshStartup(void) {
	int err;

#ifdef WIN32
	err = WSAStartup(MAKEWORD(2, 0), &_wsadata);
	if (err != 0) {
		printf("WSAStartup failed with error:  %d\n", err);
	}
#endif

	err = libssh2_init(0);
	if (err != 0) {
		printf("libssh2 initialization failed:  %d\n", err);
	}

	utilEnsureBufferSize((unsigned char **)&_output, &_outputLen, 1024);
	utilEnsureBufferSize((unsigned char **)&_buffer, &_bufferLen, 0x4000);

	g_idle_add(sshUpdate, sshUpdate);
}


gboolean sshUpdate(gpointer userData) {
	int              rc;
	int              exitcode   = 127;
	char            *exitsignal = (char *)"none";
	SSHT            *s;
	int              i;

	(void)userData;

	_output[0] = 0;

	for (i=0; i<arrlen(_activeSSHs); i++) {
		s = _activeSSHs[i];
		if (!s->finished) {
			rc = libssh2_channel_read(s->channel, _buffer, sizeof(_buffer));
			if (rc > 0) {
				utilEnsureBufferSize((unsigned char **)&_output, &_outputLen, rc + 1);
				_output[0] = 0;
				strncat(_output, (char *)_buffer, rc);
				//***TODO*** Display in status dialog
				//s->status->AddOutput(output);
			} else {
				if (rc == LIBSSH2_ERROR_EAGAIN) {
					sshWaitSocket(s->sock, s->session);
				} else {
					while ((rc = libssh2_channel_close(s->channel)) == LIBSSH2_ERROR_EAGAIN) {
						sshWaitSocket(s->sock, s->session);
					}
					DEL(_output);
					_output = strdup("\n\nUH OH!  Something wrong happened.  Check the output for errors.");
					if (rc == 0) {
						exitcode = libssh2_channel_get_exit_status(s->channel);
						libssh2_channel_get_exit_signal(s->channel, &exitsignal, NULL, NULL, NULL, NULL, NULL);
						if (exitcode == 0) {
							DEL(_output);
							_output = strdup("\n\nFINISHED!  You can close this window anytime.");
						}
					}
					//***TODO*** Display in status dialog
					//s->status->AddOutput(output);
					libssh2_channel_free(s->channel);
					s->result = exitcode;
					//s->finished = true;
					//s->status->Finished();  // No close
					if (s->callback) s->callback(s);
					arrdel(_activeSSHs, i);
					break;  // Exit for loop
					// We don't delete "s" because it will be handled by utilSSHDisconnect
				}
			}
		}
	}

	return G_SOURCE_CONTINUE;
}


static int sshWaitSocket(int socket_fd, LIBSSH2_SESSION *session) {
	struct timeval  timeout;
	int             rc;
	fd_set          fd;
	fd_set         *writefd = NULL;
	fd_set         *readfd  = NULL;
	int             dir;

	timeout.tv_sec = 10;
	timeout.tv_usec = 0;

	FD_ZERO(&fd);
	FD_SET(socket_fd, &fd);

	// Make sure we wait in the correct direction.
	dir = libssh2_session_block_directions(session);

	if (dir & LIBSSH2_SESSION_BLOCK_INBOUND) {
		readfd = &fd;
	}

	if (dir & LIBSSH2_SESSION_BLOCK_OUTBOUND) {
		writefd = &fd;
	}

	rc = select(socket_fd + 1, readfd, writefd, NULL, &timeout);

	return rc;
}