#include "stdafx.h"
#include "ProcessExecute.h"
#include "WcharMbcsConverter.h"

/* The Console Redirection is taken from the TagsView plugin from Vitaliy Dovgan.  
 * My thanks to him for pointing me in the right direction. :)
 * And of course for NppExec, without which, Notepad++ would only 
 * be half as powerful.
 */

#define DEFAULT_PIPE_SIZE 1
#define PIPE_READBUFSIZE  4096


const char *ProcessExecute::STREAM_NAME_STDOUT = "OUT";
const char *ProcessExecute::STREAM_NAME_STDERR = "ERR";

ProcessExecute::ProcessExecute()
{
}

ProcessExecute::~ProcessExecute()
{
}


bool ProcessExecute::isWindowsNT()
{
	OSVERSIONINFO osv;
	osv.dwOSVersionInfoSize = sizeof(osv);
	::GetVersionEx(&osv);
	return (osv.dwPlatformId >= VER_PLATFORM_WIN32_NT);
}

DWORD ProcessExecute::execute(const TCHAR *commandLine, boost::python::object pyStdout, boost::python::object pyStderr, boost::python::object /*pyStdin*/, bool spoolToFile /* = false */)
{
	DWORD returnValue = 0;

	if (pyStdout.is_none())
		throw process_start_exception("stdout cannot be None");
	if (pyStderr.is_none())
		throw process_start_exception("stderr cannot be None");

	// Create out, err, and in pipes (ignore in, initially)
	SECURITY_DESCRIPTOR sd;
	SECURITY_ATTRIBUTES sa;
	process_start_exception exceptionThrown("");
	
	bool thrown = false;
	
	PipeReaderArgs stdoutReaderArgs;
	PipeReaderArgs stderrReaderArgs;
	// Only used if spooling, but we need to delete it later.
	TCHAR tmpFilename[MAX_PATH] = {0};
	HANDLE hStdOutReadPipe, hStdOutWritePipe;
	HANDLE hStdErrReadPipe, hStdErrWritePipe;

	Py_BEGIN_ALLOW_THREADS
	try
	{
		if (isWindowsNT())
		{
			::InitializeSecurityDescriptor( &sd, SECURITY_DESCRIPTOR_REVISION );
			::SetSecurityDescriptorDacl( &sd, TRUE, NULL, FALSE );
			sa.lpSecurityDescriptor = &sd;
		}
		else
		{
			sa.lpSecurityDescriptor = NULL;
		}
		sa.nLength = sizeof(SECURITY_ATTRIBUTES);
		sa.bInheritHandle = TRUE;

		if (!::CreatePipe(&hStdOutReadPipe, &hStdOutWritePipe, &sa, DEFAULT_PIPE_SIZE))
		{
			throw process_start_exception("Error creating pipe for stdout");
		}

		if (!::CreatePipe(&hStdErrReadPipe, &hStdErrWritePipe, &sa, DEFAULT_PIPE_SIZE))
		{
			throw process_start_exception("Error creating pipe for stderr");
		}

		HANDLE stopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
		DWORD dwThreadId;
	

		stdoutReaderArgs.processExecute = this;
		stdoutReaderArgs.hPipeRead = hStdOutReadPipe;
		stdoutReaderArgs.hPipeWrite = hStdOutWritePipe;
		stdoutReaderArgs.pythonFile = pyStdout;
		stdoutReaderArgs.stopEvent = stopEvent;
		stdoutReaderArgs.completedEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
		stdoutReaderArgs.streamName = STREAM_NAME_STDOUT;
		stdoutReaderArgs.toFile = false;

		stderrReaderArgs.processExecute = this;
		stderrReaderArgs.hPipeRead = hStdErrReadPipe;
		stderrReaderArgs.hPipeWrite = hStdErrWritePipe;
		stderrReaderArgs.pythonFile = pyStderr;
		stderrReaderArgs.stopEvent = stopEvent;
		stderrReaderArgs.completedEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
		stderrReaderArgs.streamName = STREAM_NAME_STDERR;
		stderrReaderArgs.toFile = false;

	
	


		/* If we're in an event, we need to spool the output to a file 
		 * first, before we write it to python - so that the python writing is 
		 * done in *this* thread
		 */
		if (spoolToFile)
		{
			stdoutReaderArgs.toFile = true;
			stderrReaderArgs.toFile = true;
		
			/// Create the mutex for writing to the file
			stdoutReaderArgs.fileMutex = CreateMutex(NULL, FALSE, NULL);
			stderrReaderArgs.fileMutex = stdoutReaderArgs.fileMutex;

			/// Create the temp file
			TCHAR tmpPath[MAX_PATH];

			GetTempPath(MAX_PATH, tmpPath);
			if(!GetTempFileName(tmpPath, _T("py"), 0, tmpFilename))
			{
				throw process_start_exception("Error creating temporary filename for output spooling");
			}
			stdoutReaderArgs.file = new std::fstream(tmpFilename, std::fstream::binary | std::fstream::in | std::fstream::out);
			stderrReaderArgs.file = stdoutReaderArgs.file;
			/*
			stdoutReaderArgs.fileHandle = CreateFile(tmpFilename, GENERIC_READ | GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
			stderrReaderArgs.fileHandle = stdoutReaderArgs.fileHandle;
		
			if (INVALID_HANDLE_VALUE == stdoutReaderArgs.fileHandle)
			{
				throw process_start_exception("Error opening temporary file for output spooling");
			}
			*/


		}
		// start thread functions for stdout and stderr
		HANDLE hStdoutThread = CreateThread( 
				NULL,							// no security attribute 
				0,								// default stack size 
				pipeReader,						// thread proc
				(LPVOID) &stdoutReaderArgs,		// thread parameter 
				0,								// not suspended 
				&dwThreadId);					// returns thread ID 

		HANDLE hStderrThread = CreateThread(
				NULL,							// no security attribute 
				0,								// default stack size 
				pipeReader,						// thread proc
				(LPVOID) &stderrReaderArgs,		// thread parameter 
				0,								// not suspended 
				&dwThreadId);					// returns thread ID 


		// start process
		PROCESS_INFORMATION pi;
		STARTUPINFO         si;
	

		::SetHandleInformation(hStdOutReadPipe, HANDLE_FLAG_INHERIT, 0);
		::SetHandleInformation(hStdErrReadPipe, HANDLE_FLAG_INHERIT, 0);

		// initialize STARTUPINFO struct
		::ZeroMemory( &si, sizeof(STARTUPINFO) );
		si.cb = sizeof(STARTUPINFO);
		si.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
		si.wShowWindow = SW_HIDE;
		si.hStdInput = NULL;
		si.hStdOutput = hStdOutWritePipe;
		si.hStdError = hStdErrWritePipe;

		::ZeroMemory( &pi, sizeof(PROCESS_INFORMATION) );
	
		size_t commandLineLength = _tcslen(commandLine) + 1;
		// We use an auto_ptr here because of a potential early out due to an exception thrown.
		std::auto_ptr<TCHAR> cmdLine(new TCHAR[commandLineLength]);
		_tcscpy_s(cmdLine.get(), commandLineLength, commandLine);
		bool processStartSuccess;

		if ( ::CreateProcess(
					NULL, 
					cmdLine.get(),
					NULL,                     // security
					NULL,                     // security
					TRUE,                     // inherits handles
					CREATE_NEW_PROCESS_GROUP, // creation flags
					NULL,                     // environment
					NULL,                     // current directory
					&si,                      // startup info
					&pi                       // process info
			))
		{
			// wait for process to exit
			::WaitForSingleObject(pi.hProcess, INFINITE);
			CloseHandle(pi.hThread);	

			processStartSuccess = true;

		}
		else
		{
			processStartSuccess = false;
		}

	
	
		// Stop the reader threads
		SetEvent(stopEvent);	
		// Wait for the pipe reader threads to complete
		HANDLE handles[] = { stdoutReaderArgs.completedEvent, stderrReaderArgs.completedEvent };
		WaitForMultipleObjects(2, handles, TRUE, INFINITE);
		CloseHandle(stdoutReaderArgs.completedEvent);
		CloseHandle(stderrReaderArgs.completedEvent);
		CloseHandle(hStdoutThread);
		CloseHandle(hStderrThread);
		CloseHandle(stopEvent);
		CloseHandle(hStdOutReadPipe);
		CloseHandle(hStdOutWritePipe);
		CloseHandle(hStdErrReadPipe);
		CloseHandle(hStdErrWritePipe);

		if (processStartSuccess)
		{
			GetExitCodeProcess(pi.hProcess, &returnValue);
			CloseHandle(pi.hProcess);
		}
		else
		{
			DWORD errorNo = ::GetLastError();
			TCHAR *buffer;

			::FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_ALLOCATE_BUFFER, NULL, errorNo, 0, reinterpret_cast<LPTSTR>(&buffer), 0, NULL);

			std::shared_ptr<char> message = WcharMbcsConverter::tchar2char(buffer);
			process_start_exception ex(message.get());

			::LocalFree(buffer);
			throw ex;
	
		}
	}
	catch(process_start_exception& ex)
	{
		exceptionThrown = ex;
		thrown = true;
	}

	Py_END_ALLOW_THREADS
	
	if (spoolToFile)
	{
		CloseHandle(stdoutReaderArgs.fileMutex);
		
		spoolFile(stdoutReaderArgs.file, pyStdout, pyStderr);
		
		stdoutReaderArgs.file->close();
		if (tmpFilename[0] != 0)
		{
			DeleteFile(tmpFilename);
		}
	}

	if (thrown)
	{
		throw exceptionThrown;
	}
	
	return returnValue;
}


DWORD WINAPI ProcessExecute::pipeReader(void *args)
{
	PipeReaderArgs* pipeReaderArgs = reinterpret_cast<PipeReaderArgs*>(args);

	DWORD bytesRead;
	
	char buffer[PIPE_READBUFSIZE];
	BOOL processFinished = FALSE;

	for(;;)
	{
		::PeekNamedPipe(pipeReaderArgs->hPipeRead, NULL, 0, NULL, &bytesRead, NULL);

		if (processFinished && 0 == bytesRead)
		{
			break;
		}

		if (bytesRead > 0)
		{
			if (ReadFile(pipeReaderArgs->hPipeRead, buffer, PIPE_READBUFSIZE - 1, &bytesRead, NULL))
			{
				if (pipeReaderArgs->toFile)
				{
					pipeReaderArgs->processExecute->writeToFile(pipeReaderArgs, bytesRead, buffer);
				}
				else
				{
					pipeReaderArgs->processExecute->writeToPython(pipeReaderArgs, bytesRead, buffer);
				}
			}
		}
		else
		{
			DWORD handleIndex = WaitForSingleObject(pipeReaderArgs->stopEvent, 100);
			if (WAIT_OBJECT_0 == handleIndex)
			{
				processFinished = TRUE;
			}
		}
	}

	SetEvent(pipeReaderArgs->completedEvent);

	return 0;
}

void ProcessExecute::writeToPython(PipeReaderArgs *pipeReaderArgs, DWORD bytesRead, char *buffer)
{
	buffer[bytesRead] = '\0';
	PyGILState_STATE gstate = PyGILState_Ensure();
	try
	{
		pipeReaderArgs->pythonFile.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
	}
	catch(...)
	{
		PyErr_Print();
	}
	PyGILState_Release(gstate);
}


void ProcessExecute::writeToFile(PipeReaderArgs *pipeReaderArgs, DWORD bytesRead, char *buffer)
{
	WaitForSingleObject(pipeReaderArgs->fileMutex, INFINITE);

	pipeReaderArgs->file->write(pipeReaderArgs->streamName, ProcessExecute::STREAM_NAME_LENGTH);


	if (pipeReaderArgs->file->bad())
		throw process_start_exception("Error writing to spool file");

	char byteCount[20];
	sprintf_s(byteCount, 20, "%ul\n", bytesRead);

	pipeReaderArgs->file->write(byteCount, strlen(byteCount));

	pipeReaderArgs->file->write(buffer, bytesRead);

	
	if (pipeReaderArgs->file->bad())
		throw process_start_exception("Error writing buffer to spool file");

	ReleaseMutex(pipeReaderArgs->fileMutex);
}


void ProcessExecute::spoolFile(std::fstream* file, boost::python::object pyStdout, boost::python::object pyStderr)
{
	// Rewind
	file->seekg(0);
	char infoBuffer[30];
	char *buffer = NULL;
	size_t bufferSize = 0;
	size_t bytesToRead;

	while (!file->eof())
	{
		file->getline(infoBuffer, 30, '\n');
		if (file->eof())
			break;

		bytesToRead = strtoul(infoBuffer + STREAM_NAME_LENGTH, NULL, 0);
		if (bufferSize < bytesToRead || !buffer)
		{
			if (buffer)
				delete[] buffer;
			
			bufferSize = bytesToRead + 1;
			buffer = new char[bufferSize];
		}

		file->read(buffer, bytesToRead);
		buffer[bytesToRead] = '\0';
		if (0 == strncmp(infoBuffer, STREAM_NAME_STDOUT, STREAM_NAME_LENGTH))
		{
			// Is a stdout message
			pyStdout.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
		}
		else
		{
			// is a stderr message
			pyStderr.attr("write")(boost::python::str(const_cast<const char *>(buffer)));
		}
	}

	if (buffer)
	{
		delete[] buffer;
	}
}