/*

 * support routines for subprocess module

 *

 * Currently, this extension module is only required when using the

 * subprocess module on Windows, but in the future, stubs for other

 * platforms might be added here as well.

 *

 * Copyright (c) 2004 by Fredrik Lundh <fredrik@pythonware.com>

 * Copyright (c) 2004 by Secret Labs AB, http://www.pythonware.com

 * Copyright (c) 2004 by Peter Astrand <astrand@lysator.liu.se>

 *

 * By obtaining, using, and/or copying this software and/or its

 * associated documentation, you agree that you have read, understood,

 * and will comply with the following terms and conditions:

 *

 * Permission to use, copy, modify, and distribute this software and

 * its associated documentation for any purpose and without fee is

 * hereby granted, provided that the above copyright notice appears in

 * all copies, and that both that copyright notice and this permission

 * notice appear in supporting documentation, and that the name of the

 * authors not be used in advertising or publicity pertaining to

 * distribution of the software without specific, written prior

 * permission.

 *

 * THE AUTHORS DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,

 * INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.

 * IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY SPECIAL, INDIRECT OR

 * CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS

 * OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,

 * NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION

 * WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

 *

 */

/* Licensed to PSF under a Contributor Agreement. */

/* See http://www.python.org/2.4/license for licensing details. */

/* TODO: handle unicode command lines? */

/* TODO: handle unicode environment? */

#include "Python.h"

#define WINDOWS_LEAN_AND_MEAN

#include "windows.h"

/* -------------------------------------------------------------------- */

/* handle wrapper.  note that this library uses integers when passing

   handles to a function, and handle wrappers when returning handles.

   the wrapper is used to provide Detach and Close methods */

typedef struct {

	PyObject_HEAD

	HANDLE handle;

} sp_handle_object;

staticforward PyTypeObject sp_handle_type;

static PyObject*

sp_handle_new(HANDLE handle)

{

	sp_handle_object* self;

	self = PyObject_NEW(sp_handle_object, &sp_handle_type);

	if (self == NULL)

		return NULL;

	self->handle = handle;

	return (PyObject*) self;

}

#if defined(MS_WIN32) && !defined(MS_WIN64)

#define HANDLE_TO_PYNUM(handle)	PyInt_FromLong((long) handle)

#define PY_HANDLE_PARAM	"l"

#else

#define HANDLE_TO_PYNUM(handle)	PyLong_FromLongLong((long long) handle)

#define PY_HANDLE_PARAM	"L"

#endif

static PyObject*

sp_handle_detach(sp_handle_object* self, PyObject* args)

{

	HANDLE handle;

	if (! PyArg_ParseTuple(args, ":Detach"))

		return NULL;

	handle = self->handle;

	self->handle = NULL;

	/* note: return the current handle, as an integer */

	return HANDLE_TO_PYNUM(handle);

}

static PyObject*

sp_handle_close(sp_handle_object* self, PyObject* args)

{

	if (! PyArg_ParseTuple(args, ":Close"))

		return NULL;

	if (self->handle != INVALID_HANDLE_VALUE) {

		CloseHandle(self->handle);

		self->handle = INVALID_HANDLE_VALUE;

	}

	Py_INCREF(Py_None);

	return Py_None;

}

static void

sp_handle_dealloc(sp_handle_object* self)

{

	if (self->handle != INVALID_HANDLE_VALUE)

		CloseHandle(self->handle);

	PyObject_FREE(self);

}

static PyMethodDef sp_handle_methods[] = {

	{"Detach", (PyCFunction) sp_handle_detach, METH_VARARGS},

	{"Close",  (PyCFunction) sp_handle_close,  METH_VARARGS},

	{NULL, NULL}

};

static PyObject*

sp_handle_getattr(sp_handle_object* self, char* name)

{

	return Py_FindMethod(sp_handle_methods, (PyObject*) self, name);

}

static PyObject*

sp_handle_as_int(sp_handle_object* self)

{

	return HANDLE_TO_PYNUM(self->handle);

}

static PyNumberMethods sp_handle_as_number;

statichere PyTypeObject sp_handle_type = {

	PyObject_HEAD_INIT(NULL)

	0,				/*ob_size*/

	"_subprocess_handle", sizeof(sp_handle_object), 0,

	(destructor) sp_handle_dealloc, /*tp_dealloc*/

	0, /*tp_print*/

	(getattrfunc) sp_handle_getattr,/*tp_getattr*/

	0,				/*tp_setattr*/

	0,				/*tp_compare*/

	0,				/*tp_repr*/

	&sp_handle_as_number,		/*tp_as_number */

	0,				/*tp_as_sequence */

	0,				/*tp_as_mapping */

	0				/*tp_hash*/

};

/* -------------------------------------------------------------------- */

/* windows API functions */

static PyObject *

sp_GetStdHandle(PyObject* self, PyObject* args)

{

	HANDLE handle;

	int std_handle;

	if (! PyArg_ParseTuple(args, "i:GetStdHandle", &std_handle))

		return NULL;

	Py_BEGIN_ALLOW_THREADS

	handle = GetStdHandle((DWORD) std_handle);

	Py_END_ALLOW_THREADS

	if (handle == INVALID_HANDLE_VALUE)

		return PyErr_SetFromWindowsErr(GetLastError());

	if (! handle) {

		Py_INCREF(Py_None);

		return Py_None;

	}

	/* note: returns integer, not handle object */

	return HANDLE_TO_PYNUM(handle);

}

static PyObject *

sp_GetCurrentProcess(PyObject* self, PyObject* args)

{

	if (! PyArg_ParseTuple(args, ":GetCurrentProcess"))

		return NULL;

	return sp_handle_new(GetCurrentProcess());

}

static PyObject *

sp_DuplicateHandle(PyObject* self, PyObject* args)

{

	HANDLE target_handle;

	BOOL result;

	HANDLE source_process_handle;

	HANDLE source_handle;

	HANDLE target_process_handle;

	int desired_access;

	int inherit_handle;

	int options = 0;

	if (! PyArg_ParseTuple(args,

			       PY_HANDLE_PARAM PY_HANDLE_PARAM PY_HANDLE_PARAM

			       "ii|i:DuplicateHandle",

	                       &source_process_handle,

	                       &source_handle,

	                       &target_process_handle,

	                       &desired_access,

	                       &inherit_handle,

	                       &options))

		return NULL;

	Py_BEGIN_ALLOW_THREADS

	result = DuplicateHandle(

		source_process_handle,

		source_handle,

		target_process_handle,

		&target_handle,

		desired_access,

		inherit_handle,

		options

	);

	Py_END_ALLOW_THREADS

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());

	return sp_handle_new(target_handle);

}

static PyObject *

sp_CreatePipe(PyObject* self, PyObject* args)

{

	HANDLE read_pipe;

	HANDLE write_pipe;

	BOOL result;

	PyObject* pipe_attributes; /* ignored */

	int size;

	if (! PyArg_ParseTuple(args, "Oi:CreatePipe", &pipe_attributes, &size))

		return NULL;

	Py_BEGIN_ALLOW_THREADS

	result = CreatePipe(&read_pipe, &write_pipe, NULL, size);

	Py_END_ALLOW_THREADS

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());

	return Py_BuildValue(

		"NN", sp_handle_new(read_pipe), sp_handle_new(write_pipe));

}

/* helpers for createprocess */

static int

getint(PyObject* obj, char* name)

{

	PyObject* value;

	int ret;

	value = PyObject_GetAttrString(obj, name);

	if (! value) {

		PyErr_Clear(); /* FIXME: propagate error? */

		return 0;

	}

	ret = (int) PyInt_AsLong(value);

	Py_DECREF(value);

	return ret;

}

static HANDLE

gethandle(PyObject* obj, char* name)

{

	sp_handle_object* value;

	HANDLE ret;

	value = (sp_handle_object*) PyObject_GetAttrString(obj, name);

	if (! value) {

		PyErr_Clear(); /* FIXME: propagate error? */

		return NULL;

	}

	if (value->ob_type != &sp_handle_type)

		ret = NULL;

	else

		ret = value->handle;

	Py_DECREF(value);

	return ret;

}

static PyObject*

getenvironment(PyObject* environment)

{

	int i, envsize;

	PyObject* out = NULL;

	PyObject* keys;

	PyObject* values;

	char* p;

	/* convert environment dictionary to windows enviroment string */

	if (! PyMapping_Check(environment)) {

		PyErr_SetString(

		    PyExc_TypeError, "environment must be dictionary or None");

		return NULL;

	}

	envsize = PyMapping_Length(environment);

	keys = PyMapping_Keys(environment);

	values = PyMapping_Values(environment);

	if (!keys || !values)

		goto error;

	out = PyString_FromStringAndSize(NULL, 2048);

	if (! out)

		goto error;

	p = PyString_AS_STRING(out);

	for (i = 0; i < envsize; i++) {

		int ksize, vsize, totalsize;

		PyObject* key = PyList_GET_ITEM(keys, i);

		PyObject* value = PyList_GET_ITEM(values, i);

		if (! PyString_Check(key) || ! PyString_Check(value)) {

			PyErr_SetString(PyExc_TypeError,

				"environment can only contain strings");

			goto error;

		}

		ksize = PyString_GET_SIZE(key);

		vsize = PyString_GET_SIZE(value);

		totalsize = (p - PyString_AS_STRING(out)) + ksize + 1 +

							     vsize + 1 + 1;

		if (totalsize > PyString_GET_SIZE(out)) {

			int offset = p - PyString_AS_STRING(out);

			_PyString_Resize(&out, totalsize + 1024);

			p = PyString_AS_STRING(out) + offset;

		}

		memcpy(p, PyString_AS_STRING(key), ksize);

		p += ksize;

		*p++ = '=';

		memcpy(p, PyString_AS_STRING(value), vsize);

		p += vsize;

		*p++ = '\0';

	}

	/* add trailing null byte */

	*p++ = '\0';

	_PyString_Resize(&out, p - PyString_AS_STRING(out));

	/* PyObject_Print(out, stdout, 0); */

	Py_XDECREF(keys);

	Py_XDECREF(values);

	return out;

 error:

	Py_XDECREF(out);

	Py_XDECREF(keys);

	Py_XDECREF(values);

	return NULL;

}

static PyObject *

sp_CreateProcess(PyObject* self, PyObject* args)

{

	BOOL result;

	PROCESS_INFORMATION pi;

	STARTUPINFO si;

	PyObject* environment;

	char* application_name;

	char* command_line;

	PyObject* process_attributes; /* ignored */

	PyObject* thread_attributes; /* ignored */

	int inherit_handles;

	int creation_flags;

	PyObject* env_mapping;

	char* current_directory;

	PyObject* startup_info;

	if (! PyArg_ParseTuple(args, "zzOOiiOzO:CreateProcess",

			       &application_name,

			       &command_line,

			       &process_attributes,

			       &thread_attributes,

			       &inherit_handles,

			       &creation_flags,

			       &env_mapping,

			       &current_directory,

			       &startup_info))

		return NULL;

	ZeroMemory(&si, sizeof(si));

	si.cb = sizeof(si);

	/* note: we only support a small subset of all SI attributes */

	si.dwFlags = getint(startup_info, "dwFlags");

	si.wShowWindow = getint(startup_info, "wShowWindow");

	si.hStdInput = gethandle(startup_info, "hStdInput");

	si.hStdOutput = gethandle(startup_info, "hStdOutput");

	si.hStdError = gethandle(startup_info, "hStdError");

	if (PyErr_Occurred())

		return NULL;

	if (env_mapping == Py_None)

		environment = NULL;

	else {

		environment = getenvironment(env_mapping);

		if (! environment)

			return NULL;

	}

	Py_BEGIN_ALLOW_THREADS

	result = CreateProcess(application_name,

			       command_line,

			       NULL,

			       NULL,

			       inherit_handles,

			       creation_flags,

			       environment ? PyString_AS_STRING(environment) : NULL,

			       current_directory,

			       &si,

			       &pi);

	Py_END_ALLOW_THREADS

	Py_XDECREF(environment);

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());

	return Py_BuildValue("NNii",

			     sp_handle_new(pi.hProcess),

			     sp_handle_new(pi.hThread),

			     pi.dwProcessId,

			     pi.dwThreadId);

}

static PyObject *

sp_TerminateProcess(PyObject* self, PyObject* args)

{

	BOOL result;

	HANDLE process;

	int exit_code;

	if (! PyArg_ParseTuple(args, PY_HANDLE_PARAM "i:TerminateProcess",

			       &process, &exit_code))

		return NULL;

	result = TerminateProcess(process, exit_code);

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());

	Py_INCREF(Py_None);

	return Py_None;

}

static PyObject *

sp_GetExitCodeProcess(PyObject* self, PyObject* args)

{

	DWORD exit_code;

	BOOL result;

	HANDLE process;

	if (! PyArg_ParseTuple(args, PY_HANDLE_PARAM ":GetExitCodeProcess", &process))

		return NULL;

	result = GetExitCodeProcess(process, &exit_code);

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());

	return PyInt_FromLong(exit_code);

}

static PyObject *

sp_WaitForSingleObject(PyObject* self, PyObject* args)

{

	DWORD result;

	HANDLE handle;

	int milliseconds;

	if (! PyArg_ParseTuple(args, PY_HANDLE_PARAM "i:WaitForSingleObject",

	                  	     &handle,

	                  	     &milliseconds))

		return NULL;

	Py_BEGIN_ALLOW_THREADS

	result = WaitForSingleObject(handle, (DWORD) milliseconds);

	Py_END_ALLOW_THREADS

	if (result == WAIT_FAILED)

		return PyErr_SetFromWindowsErr(GetLastError());

	return PyInt_FromLong((int) result);

}

static PyObject *

sp_GetVersion(PyObject* self, PyObject* args)

{

	if (! PyArg_ParseTuple(args, ":GetVersion"))

		return NULL;

	return PyInt_FromLong((int) GetVersion());

}

static PyObject *

sp_GetModuleFileName(PyObject* self, PyObject* args)

{

	BOOL result;

	HMODULE module;

	TCHAR filename[MAX_PATH];

	if (! PyArg_ParseTuple(args, PY_HANDLE_PARAM ":GetModuleFileName",

			       &module))

		return NULL;

	result = GetModuleFileName(module, filename, MAX_PATH);

	filename[MAX_PATH-1] = '\0';

	if (! result)

		return PyErr_SetFromWindowsErr(GetLastError());