/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2012 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <mitsuba/core/sshstream.h>
#include <mitsuba/core/statistics.h>

#if !defined(WIN32)
# include <unistd.h>
#else
# include <windows.h>
#endif
#include <errno.h>

MTS_NAMESPACE_BEGIN

struct SSHStream::SSHStreamPrivate
{
	const std::string userName, hostName;
	const int port, timeout;
	size_t received, sent;
#if defined(WIN32)
	HANDLE childInRd,  childInWr;
	HANDLE childOutRd, childOutWr;
#else
	int infd, outfd;
	FILE *input, *output;
#endif

	SSHStreamPrivate(const std::string& uname, const std::string& hname,
		int p, int tm) :
	userName(uname), hostName(hname), port(p), timeout(tm), received(0), sent(0),
#if defined(WIN32)
	childInRd(0), childInWr(0), childOutRd(0), childOutWr(0)
#else
	infd(-1), outfd(-1), input(0), output(0)
#endif
	{}
};

SSHStream::SSHStream(const std::string &userName,
	const std::string &hostName, const std::vector<std::string> &cmdLine, int port, int timeout)
 : d(new SSHStreamPrivate(userName, hostName, port, timeout)) {
	setByteOrder(ENetworkByteOrder);

	Log(EInfo, "Establishing a SSH connection to \"%s@%s\"", 
		userName.c_str(), hostName.c_str());

#if defined(WIN32)
	/* Inherit pipe handles */
	SECURITY_ATTRIBUTES sAttr;
	sAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
	sAttr.bInheritHandle = TRUE;
	sAttr.lpSecurityDescriptor = NULL;

	/* Create stdout pipe */
	if (!CreatePipe(&(d->childOutRd), &(d->childOutWr), &sAttr, 0))
		Log(EError, "Error in CreatePipe(): %s", lastErrorText().c_str());

	/* Create stdin pipe */
	if (!CreatePipe(&(d->childInRd), &(d->childInWr), &sAttr, 0))
		Log(EError, "Error in CreatePipe(): %s", lastErrorText().c_str());

	/* Only inherit one side of the pipes */
	if (!SetHandleInformation(d->childOutRd, HANDLE_FLAG_INHERIT, 0))
		Log(EError, "Error in SetHandleInformation(): %s", lastErrorText().c_str());
	if (!SetHandleInformation(d->childInWr, HANDLE_FLAG_INHERIT, 0))
		Log(EError, "Error in SetHandleInformation(): %s", lastErrorText().c_str());

	/* Start the plink process */
	PROCESS_INFORMATION pi;
	STARTUPINFO si;
	ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));
	ZeroMemory(&si, sizeof(STARTUPINFO));
	si.cb = sizeof(STARTUPINFO);
	si.hStdError  = d->childOutWr;
	si.hStdOutput = d->childOutWr;
	si.hStdInput  = d->childInRd;
	si.dwFlags |= STARTF_USESTDHANDLES;

	std::string params = formatString("-batch -P %i -T %s@%s",
		d->port, d->userName.c_str(), d->hostName.c_str());

	for (size_t i=0; i<cmdLine.size(); ++i)
		params = params + " " + cmdLine[i];

	if (!CreateProcess("plink.exe", (LPSTR) params.c_str(), 
		NULL, NULL, // Process & thread security attributes
		TRUE, // Inherit handles
		CREATE_NEW_PROCESS_GROUP, // Ignore Ctrl-C
		NULL, // Use environment of the calling process
		NULL, // Use CWD of the calling process
		&si, &pi))
		Log(EError, "Are you sure plink.exe exists? Take a look at sshstream.h "
			"-- OS error: CreateProcess() failed: %s", lastErrorText().c_str());

	CloseHandle(pi.hProcess);
	CloseHandle(pi.hThread);
	CloseHandle(d->childOutWr);
	CloseHandle(d->childInRd);
#else
	int infd[2], outfd[2];

	if (pipe(outfd) == -1)
		Log(EError, "Error in pipe(): %s", strerror(errno));

	if (pipe(infd) == -1)
		Log(EError, "Error in pipe(): %s", strerror(errno));

	char **argv = new char*[15+cmdLine.size()];
	int argc = 0;

	argv[argc++] = strdup("ssh");
	if (d->port != 22) {
		argv[argc++] = strdup("-p"); argv[argc++] = strdup(formatString("%i", d->port).c_str());
	}
	argv[argc++] = strdup("-e"); argv[argc++] = strdup("none"); /* No escape character */
	argv[argc++] = strdup("-T"); /* Disable pseudo TTY allocation (need to be able to transfer binary data) */
	argv[argc++] = strdup("-o"); argv[argc++] = strdup("PasswordAuthentication no"); /* Don't proceed if password auth. is required */
	argv[argc++] = strdup("-o"); argv[argc++] = strdup("StrictHostKeyChecking no"); /* Don't ask whether to accept host keys */
	argv[argc++] = strdup("-o"); argv[argc++] = strdup(formatString("ConnectTimeout %i", d->timeout).c_str()); /* Don't get stuck too long */
	argv[argc++] = strdup("-q"); /* Quiet */
	argv[argc++] = strdup(formatString("%s@%s", d->userName.c_str(), d->hostName.c_str()).c_str());
	for (size_t i=0; i<cmdLine.size(); ++i)
		argv[argc++] = strdup(cmdLine[i].c_str());
	argv[argc++] = NULL;

	if (!fork()) {
		if (close(0) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (close(1) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (dup2(outfd[0], 0) == -1)
			Log(EError, "Error in dup2(): %s!", strerror(errno));
		if (dup2(infd[1], 1) == -1)
			Log(EError, "Error in dup2(): %s!", strerror(errno));
		if (close(outfd[0]) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (close(outfd[1]) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (close(infd[0]) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (close(infd[1]) == -1)
			Log(EError, "Error in close(): %s!", strerror(errno));
		if (execvp("/usr/bin/ssh", (char * const *) argv) == -1)
			Log(EError, "Error in execvp(): %s!", strerror(errno));
	} else {
		close(outfd[0]);
		close(infd[1]);
		d->infd = infd[0];
		d->outfd = outfd[1];
		d->input = fdopen(infd[0], "rb");
		d->output = fdopen(outfd[1], "wb");
	}
	for (int i=0; i<argc-1; ++i)
		free(argv[i]);
#endif
}

SSHStream::~SSHStream() {
	Log(EDebug, "Closing SSH connection");
#if defined(WIN32)
	CloseHandle(d->childInWr);
	CloseHandle(d->childOutRd);
#else
	fclose(d->input);
	fclose(d->output);
#endif
}

const std::string& SSHStream::getHostName() const {
	return d->hostName;
}
	
const std::string& SSHStream::getUserName() const {
	return d->userName;
}
	
size_t SSHStream::getReceivedBytes() const {
	return d->received;
}
	
size_t SSHStream::getSentBytes() const {
	return d->sent;
}

std::string SSHStream::toString() const {
	std::ostringstream oss;
	oss << "SSHStream[userName='"<< d->userName << "', hostName='" 
		<< d->hostName << "', sent=" << (d->sent / 1024) << " KB, "
		"received=" << (d->received/1024) << " KB]" << endl;
	return oss.str();
}

void SSHStream::seek(size_t pos) {
	Log(EError, "Cannot seek within a socket stream!");
}

size_t SSHStream::getPos() const {
	Log(EError, "Cannot determine the position within a socket stream!");
	return 0;
}

size_t SSHStream::getSize() const {
	Log(EError, "Cannot determine the size of a socket stream!");
	return 0;
}

void SSHStream::truncate(size_t size) {
	Log(EError, "Cannot truncate a socket stream!");
}

void SSHStream::flush() {
#if defined(WIN32)
	// No-op
#else
	if (fflush(d->output) == EOF)
		Log(EError, "Error in fflush(): %s!", strerror(errno));
#endif
}

void SSHStream::read(void *ptr, size_t size) {
	static StatsCounter bytesRcvd("Network", "Bytes received (SSH)");
#if defined(WIN32)
	size_t left = size;
	char *data = (char *) ptr;
	while (left > 0) {
		DWORD nRead = 0;
		if (!ReadFile(d->childOutRd, ptr, (DWORD) left, &nRead, NULL))
			Log(EError, "Connection closed while reading: %s", lastErrorText().c_str());
		left -= nRead;
		data += nRead;
	}
#else
	if (fread(ptr, size, 1, d->input) != 1) {
		if (feof(d->input))
			Log(EError, "Error in fread(): end of file!");
		else if (ferror(d->input))
			Log(EError, "Error in fread(): stream error!");
		/* Otherwise, ignore (strange, but seems to be required) */
	}
#endif
	d->received += size;
	bytesRcvd += size;
}

void SSHStream::write(const void *ptr, size_t size) {
	static StatsCounter bytesSent("Network", "Bytes sent (SSH)");
#if defined(WIN32)
	size_t left = size;
	char *data = (char *) ptr;
	while (left > 0) {
		DWORD nWritten = 0;
		if (!WriteFile(d->childInWr, ptr, (DWORD) left, &nWritten, NULL))
			Log(EError, "Connection closed while writing: %s", lastErrorText().c_str());
		left -= nWritten;
		data += nWritten;
	}
#else
	if (fwrite(ptr, size, 1, d->output) != 1) {
		if (feof(d->output))
			Log(EError, "Error in fwrite(): end of file!");
		else if (ferror(d->output))
			Log(EError, "Error in fwrite(): stream error!");
		/* Otherwise, ignore (strange, but seems to be required) */
	}
#endif
	d->sent += size;
	bytesSent += size;
}

bool SSHStream::canRead() const {
	return true;
}

bool SSHStream::canWrite() const {
	return true;
}

MTS_IMPLEMENT_CLASS(SSHStream, false, Stream)
MTS_NAMESPACE_END