/*
** SellaNMS
** $Id: ssh.cxx,v 1.4 2006/01/24 07:08:09 sella Exp $
** Copyright (c) 2001-2005 Digital Genesis Software, LLC. All Rights Reserved.
** Released under the GPL Version 2 License.
** http://www.digitalgenesis.com
*/

static const char rcsid[] = "$Id: ssh.cxx,v 1.4 2006/01/24 07:08:09 sella Exp $";

#include "ssh.h"
#include "util.h"

#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/poll.h>
#include <unistd.h>
#include <regex.h>
#include <errno.h>
#include <utmp.h>
#include <pty.h>
#include <stdio.h>

#include <stl/list>
#include <stl/slist>
#include <stl/string>

extern int errno;

SSH::SSH():ptty(0) {
	string	regex;

	/* All of these regular expressions are OpenSSH version dependant.  Other versions of OpenSSH
		or brands of SSH client may not work. Something should be done for compatability, but
		since I only need to work with my version of SSH, it's good enough for me. */
	regex = "The authenticity of host '.*' can't be established..*Are you sure you want to continue connecting (yes/no)?";
	if (regcomp(&(this->authRegex), regex.c_str(), REG_EXTENDED) != 0) {
		fprintf(stderr, "Failed to compile regular expression.\n");
	}

	regex = "^.*@.*'s password: $";
	if (regcomp(&(this->passwordRegex), regex.c_str(), REG_EXTENDED) != 0) {
		fprintf(stderr, "Failed to compile regular expression.\n");
	}

	regex = "^Please type 'yes' or 'no': $";
	if (regcomp(&(this->yesnoRegex), regex.c_str(), REG_EXTENDED) != 0) {
		fprintf(stderr, "Failed to compile regular expression.\n");
	}
}

SSH::~SSH() {
	this->disconnect();

	regfree(&(this->authRegex));
	regfree(&(this->passwordRegex));
	regfree(&(this->yesnoRegex));
	regfree(&(this->shellRegex));
}

bool SSH::connect(string hostname, string username, string password, string prompt) {
	char			buffer[10024];
	bool			connected = false;

	size_t		match_count = 1;
	regmatch_t	matches;

	pid_t			child_process;
	string		transmission;

	child_process = forkpty(&(this->ptty), NULL, NULL, NULL);
	if (child_process < 0) {
		/* An error occured while attempting to fork this process. */
		fprintf(stderr, "Failed to fork process for SSH connection: %s.\n", strerror(errno));

		return false;
	} else if (child_process == 0) {
		/* Replace this process image with the system command ssh. */
		execlp("ssh", "ssh", "-t", "-l", username.c_str(), hostname.c_str(), NULL);

		/* There was an erorr executing SSH: execlp should never return. */
		fprintf(stderr, "Failed to execute the SSH command while attempting to connect to %s.", hostname.c_str());
		return false;
	} else {
		while (!connected) {
			usleep(10000);

			/* Read system output from PTTY. */
			this->recieve(transmission);

			/* Look for password prompt. */
			if (regexec(&(this->passwordRegex), transmission.c_str(), match_count, &matches, 0) != REG_NOMATCH) {
				/* Send password to host. */
				transmission = password + "\r";
				this->transmit(transmission);

				connected = true;
			}

			/* Look for first time permission request prompt. */
			if (!regexec(&(this->authRegex), transmission.c_str(), match_count, &matches, 0)) {
				/* Send yes to host. */
				transmission = "yes\r";
				this->transmit(transmission);
			}

			/* Look for yes/no prompt. */
			if (!regexec(&(this->yesnoRegex), transmission.c_str(), match_count, &matches, 0)) {
				/* Send yes to host. */
				transmission = "yes\r";
				this->transmit(transmission);
			}
		}
		connected = false;

		if (regcomp(&(this->shellRegex), prompt.c_str(), REG_EXTENDED) != 0) {
			fprintf(stderr, "Failed to compile shell prompt regular expression.\n");
		}
		while (!connected) {
			/* Read system output from PTTY. */
			this->recieve(transmission);

			/* Look for shell prompt. */
			if (!regexec(&(this->shellRegex), transmission.c_str(), match_count, &matches, 0)) {
				connected = true;
			} else {
				usleep(10000);
			}
		}
	}
	return true;
}

bool SSH::disconnect() {
	if (this->ptty != 0) {
		close(this->ptty);

		this->ptty = 0;
	}
	return true;
}

bool SSH::recieve(string &value) {
	char		buffer[1024];
	size_t	count = 1023;
	
	pollfd	pfd[1];
	
	pfd[0].fd = this->ptty;	
	pfd[0].events = POLLIN | POLLPRI;

	value.erase();
	memset(buffer, 0, 1024);

	/* Do not use select here instead of poll.  It appears to result in a race condition
		which can hang read() indefinately. This has to do with (this->ptty) being a full
		duplex file descriptor. */
	while (poll(pfd, 1, 100) > 0) {
		count = read(this->ptty, buffer, 1023);

		if (count < 0) {
			fprintf(stderr, "Failed to read from ssh session: %s\n", strerror(errno));
		} else if (count == 0) {
			break;
		}

		value.append(buffer);
		memset(buffer, 0, 1024);
	}
	return true;
}

bool SSH::transmit(string &value) {
	pollfd	pfd[1];
	
	pfd[0].fd = this->ptty;	
	pfd[0].events = POLLOUT;

	/* Do not use select here instead of poll.  It appears to result in a race condition
		which can hang write() indefinately. This has to do with (this->ptty) being a full
		duplex file descriptor. */
	while (poll(pfd, 1, 100) < 1) {
		usleep(0);
	}
	if (write(this->ptty, value.c_str(), value.size()) < 0) {
		fprintf(stderr, "Failure transmitting over SSH tunnel: %s\n", strerror(errno));

		return false;
	}
	value.erase();
	return true;
}

bool SSH::sendCommand(string const &command) {
	string	cmd = command + "\r";

	return this->transmit(cmd);
}

bool SSH::getReply(string &reply) {
	return this->recieve(reply);
}

bool SSH::getReply(list<string> &reply) {
	string	bulk;
	string	line;

	bool		escape = false;

	reply.clear();
	if (!this->recieve(bulk)) {
		return false;
	}

	/* Segment the result into lines and strip terminal escape sequences (color codes, etc...). */
	for (string::size_type i = 0; i < bulk.size(); i++) {
		switch (bulk[i]) {
			case 0x1b:	// Escape
				escape = true;
				break;

			case '\r':
			case '\n':
				if (line.size() > 0) {
					reply.push_back(line);

					line.erase();
				}
				break;

			case 'm':	// Comes at end of escape
				if (escape) {
					escape = false;

					break;
				}
			default:
				if (!escape) {
					line.push_back(bulk[i]);
				}
		}
	}
	return true;
}

bool SSH::sendCommand(string const &command, list<string> &result) {
	bool				prompt = false;
	size_t			match_count = 1;

	regmatch_t		matches;

	list<string>	verify;
	string			partial;

	result.clear();
	if (!this->sendCommand(command)) {
		return false;
	}
	if (!this->getReply(result)) {
		return false;
	}

	/* Erase the first item in the result since it is our command. */
	result.pop_front();

	while (!prompt) {
		usleep(10000);

		/* Look for shell prompt. */
		if (!regexec(&(this->shellRegex), result.back().c_str(), match_count, &matches, 0)) {
			/* Make sure that we didn't just get a match back on the result of some command
				that has a command prompt in it's result. */
			if (this->getReply(verify)) {
				for (list<string>::iterator s = verify.begin(); s != verify.end(); ++s) {
					result.push_back(*s);
				}
				if (verify.size() == 0) {
					prompt = true;
				}
				verify.clear();
			}
		} else {
			/* No shell prompt.  This must not be the end of output yet. */
			if (this->getReply(verify)) {
				for (list<string>::iterator s = verify.begin(); s != verify.end(); ++s) {
					result.push_back(*s);
				}
				verify.clear();
			}
		}
	}


	return true;
}

/*
** Local Variables:
** c-basic-offset: 3
** tab-width: 3
** End:
** vim:
** vim: noet ts=3 sw=3
*/
