1 /*

     2  * support routines for subprocess module

     3  *

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

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

     6  * platforms might be added here as well.

     7  *

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

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

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

    11  *

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

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

    14  * and will comply with the following terms and conditions:

    15  *

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

    17  * its associated documentation for any purpose and without fee is

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

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

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

    21  * authors not be used in advertising or publicity pertaining to

    22  * distribution of the software without specific, written prior

    23  * permission.

    24  *

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

    26  * INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.

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

    28  * CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS

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

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

    31  * WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

    32  *

    33  */

    34 

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

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

    37 

    38 /* TODO: handle unicode command lines? */

    39 /* TODO: handle unicode environment? */

    40 

    41 #include "Python.h"

    42 

    43 #define WINDOWS_LEAN_AND_MEAN

    44 #include "windows.h"

    45 

    46 /* -------------------------------------------------------------------- */

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

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

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

    50 

    51 typedef struct {

    52 	PyObject_HEAD

    53 	HANDLE handle;

    54 } sp_handle_object;

    55 

    56 staticforward PyTypeObject sp_handle_type;

    57 

    58 static PyObject*

    59 sp_handle_new(HANDLE handle)

    60 {

    61 	sp_handle_object* self;

    62 

    63 	self = PyObject_NEW(sp_handle_object, &sp_handle_type);

    64 	if (self == NULL)

    65 		return NULL;

    66 

    67 	self->handle = handle;

    68 

    69 	return (PyObject*) self;

    70 }

    71 

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

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

    74 #define PY_HANDLE_PARAM	"l"

    75 #else

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

    77 #define PY_HANDLE_PARAM	"L"

    78 #endif

    79 

    80 static PyObject*

    81 sp_handle_detach(sp_handle_object* self, PyObject* args)

    82 {

    83 	HANDLE handle;

    84 

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

    86 		return NULL;

    87 

    88 	handle = self->handle;

    89 

    90 	self->handle = NULL;

    91 

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

    93 	return HANDLE_TO_PYNUM(handle);

    94 }

    95 

    96 static PyObject*

    97 sp_handle_close(sp_handle_object* self, PyObject* args)

    98 {

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

   100 		return NULL;

   101 

   102 	if (self->handle != INVALID_HANDLE_VALUE) {

   103 		CloseHandle(self->handle);

   104 		self->handle = INVALID_HANDLE_VALUE;

   105 	}

   106 	Py_INCREF(Py_None);

   107 	return Py_None;

   108 }

   109 

   110 static void

   111 sp_handle_dealloc(sp_handle_object* self)

   112 {

   113 	if (self->handle != INVALID_HANDLE_VALUE)

   114 		CloseHandle(self->handle);

   115 	PyObject_FREE(self);

   116 }

   117 

   118 static PyMethodDef sp_handle_methods[] = {

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

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

   121 	{NULL, NULL}

   122 };

   123 

   124 static PyObject*

   125 sp_handle_getattr(sp_handle_object* self, char* name)

   126 {

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

   128 }

   129 

   130 static PyObject*

   131 sp_handle_as_int(sp_handle_object* self)

   132 {

   133 	return HANDLE_TO_PYNUM(self->handle);

   134 }

   135 

   136 static PyNumberMethods sp_handle_as_number;

   137 

   138 statichere PyTypeObject sp_handle_type = {

   139 	PyObject_HEAD_INIT(NULL)

   140 	0,				/*ob_size*/

   141 	"_subprocess_handle", sizeof(sp_handle_object), 0,

   142 	(destructor) sp_handle_dealloc, /*tp_dealloc*/

   143 	0, /*tp_print*/

   144 	(getattrfunc) sp_handle_getattr,/*tp_getattr*/

   145 	0,				/*tp_setattr*/

   146 	0,				/*tp_compare*/

   147 	0,				/*tp_repr*/

   148 	&sp_handle_as_number,		/*tp_as_number */

   149 	0,				/*tp_as_sequence */

   150 	0,				/*tp_as_mapping */

   151 	0				/*tp_hash*/

   152 };

   153 

   154 /* -------------------------------------------------------------------- */

   155 /* windows API functions */

   156 

   157 static PyObject *

   158 sp_GetStdHandle(PyObject* self, PyObject* args)

   159 {

   160 	HANDLE handle;

   161 	int std_handle;

   162 

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

   164 		return NULL;

   165 

   166 	Py_BEGIN_ALLOW_THREADS

   167 	handle = GetStdHandle((DWORD) std_handle);

   168 	Py_END_ALLOW_THREADS

   169 

   170 	if (handle == INVALID_HANDLE_VALUE)

   171 		return PyErr_SetFromWindowsErr(GetLastError());

   172 

   173 	if (! handle) {

   174 		Py_INCREF(Py_None);

   175 		return Py_None;

   176 	}

   177 

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

   179 	return HANDLE_TO_PYNUM(handle);

   180 }

   181 

   182 static PyObject *

   183 sp_GetCurrentProcess(PyObject* self, PyObject* args)

   184 {

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

   186 		return NULL;

   187 

   188 	return sp_handle_new(GetCurrentProcess());

   189 }

   190 

   191 static PyObject *

   192 sp_DuplicateHandle(PyObject* self, PyObject* args)

   193 {

   194 	HANDLE target_handle;

   195 	BOOL result;

   196 

   197 	HANDLE source_process_handle;

   198 	HANDLE source_handle;

   199 	HANDLE target_process_handle;

   200 	int desired_access;

   201 	int inherit_handle;

   202 	int options = 0;

   203 

   204 	if (! PyArg_ParseTuple(args,

   205 			       PY_HANDLE_PARAM PY_HANDLE_PARAM PY_HANDLE_PARAM

   206 			       "ii|i:DuplicateHandle",

   207 	                       &source_process_handle,

   208 	                       &source_handle,

   209 	                       &target_process_handle,

   210 	                       &desired_access,

   211 	                       &inherit_handle,

   212 	                       &options))

   213 		return NULL;

   214 

   215 	Py_BEGIN_ALLOW_THREADS

   216 	result = DuplicateHandle(

   217 		source_process_handle,

   218 		source_handle,

   219 		target_process_handle,

   220 		&target_handle,

   221 		desired_access,

   222 		inherit_handle,

   223 		options

   224 	);

   225 	Py_END_ALLOW_THREADS

   226 

   227 	if (! result)

   228 		return PyErr_SetFromWindowsErr(GetLastError());

   229 

   230 	return sp_handle_new(target_handle);

   231 }

   232 

   233 static PyObject *

   234 sp_CreatePipe(PyObject* self, PyObject* args)

   235 {

   236 	HANDLE read_pipe;

   237 	HANDLE write_pipe;

   238 	BOOL result;

   239 

   240 	PyObject* pipe_attributes; /* ignored */

   241 	int size;

   242 

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

   244 		return NULL;

   245 

   246 	Py_BEGIN_ALLOW_THREADS

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

   248 	Py_END_ALLOW_THREADS

   249 

   250 	if (! result)

   251 		return PyErr_SetFromWindowsErr(GetLastError());

   252 

   253 	return Py_BuildValue(

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

   255 }

   256 

   257 /* helpers for createprocess */

   258 

   259 static int

   260 getint(PyObject* obj, char* name)

   261 {

   262 	PyObject* value;

   263 	int ret;

   264 

   265 	value = PyObject_GetAttrString(obj, name);

   266 	if (! value) {

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

   268 		return 0;

   269 	}

   270 	ret = (int) PyInt_AsLong(value);

   271 	Py_DECREF(value);

   272 	return ret;

   273 }

   274 

   275 static HANDLE

   276 gethandle(PyObject* obj, char* name)

   277 {

   278 	sp_handle_object* value;

   279 	HANDLE ret;

   280 

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

   282 	if (! value) {

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

   284 		return NULL;

   285 	}

   286 	if (value->ob_type != &sp_handle_type)

   287 		ret = NULL;

   288 	else

   289 		ret = value->handle;

   290 	Py_DECREF(value);

   291 	return ret;

   292 }

   293 

   294 static PyObject*

   295 getenvironment(PyObject* environment)

   296 {

   297 	int i, envsize;

   298 	PyObject* out = NULL;

   299 	PyObject* keys;

   300 	PyObject* values;

   301 	char* p;

   302 

   303 	/* convert environment dictionary to windows enviroment string */

   304 	if (! PyMapping_Check(environment)) {

   305 		PyErr_SetString(

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

   307 		return NULL;

   308 	}

   309 

   310 	envsize = PyMapping_Length(environment);

   311 

   312 	keys = PyMapping_Keys(environment);

   313 	values = PyMapping_Values(environment);

   314 	if (!keys || !values)

   315 		goto error;

   316 

   317 	out = PyString_FromStringAndSize(NULL, 2048);

   318 	if (! out)

   319 		goto error;

   320 

   321 	p = PyString_AS_STRING(out);

   322 

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

   324 		int ksize, vsize, totalsize;

   325 		PyObject* key = PyList_GET_ITEM(keys, i);

   326 		PyObject* value = PyList_GET_ITEM(values, i);

   327 

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

   329 			PyErr_SetString(PyExc_TypeError,

   330 				"environment can only contain strings");

   331 			goto error;

   332 		}

   333 		ksize = PyString_GET_SIZE(key);

   334 		vsize = PyString_GET_SIZE(value);

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

   336 							     vsize + 1 + 1;

   337 		if (totalsize > PyString_GET_SIZE(out)) {

   338 			int offset = p - PyString_AS_STRING(out);

   339 			_PyString_Resize(&out, totalsize + 1024);

   340 			p = PyString_AS_STRING(out) + offset;

   341 		}

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

   343 		p += ksize;

   344 		*p++ = '=';

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

   346 		p += vsize;

   347 		*p++ = '\0';

   348 	}

   349 

   350 	/* add trailing null byte */

   351 	*p++ = '\0';

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

   353 

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

   355 

   356 	Py_XDECREF(keys);

   357 	Py_XDECREF(values);

   358 

   359 	return out;

   360 

   361  error:

   362 	Py_XDECREF(out);

   363 	Py_XDECREF(keys);

   364 	Py_XDECREF(values);

   365 	return NULL;

   366 }

   367 

   368 static PyObject *

   369 sp_CreateProcess(PyObject* self, PyObject* args)

   370 {

   371 	BOOL result;

   372 	PROCESS_INFORMATION pi;

   373 	STARTUPINFO si;

   374 	PyObject* environment;

   375 

   376 	char* application_name;

   377 	char* command_line;

   378 	PyObject* process_attributes; /* ignored */

   379 	PyObject* thread_attributes; /* ignored */

   380 	int inherit_handles;

   381 	int creation_flags;

   382 	PyObject* env_mapping;

   383 	char* current_directory;

   384 	PyObject* startup_info;

   385 

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

   387 			       &application_name,

   388 			       &command_line,

   389 			       &process_attributes,

   390 			       &thread_attributes,

   391 			       &inherit_handles,

   392 			       &creation_flags,

   393 			       &env_mapping,

   394 			       &current_directory,

   395 			       &startup_info))

   396 		return NULL;

   397 

   398 	ZeroMemory(&si, sizeof(si));

   399 	si.cb = sizeof(si);

   400 

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

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

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

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

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

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

   407 

   408 	if (PyErr_Occurred())

   409 		return NULL;

   410 

   411 	if (env_mapping == Py_None)

   412 		environment = NULL;

   413 	else {

   414 		environment = getenvironment(env_mapping);

   415 		if (! environment)

   416 			return NULL;

   417 	}

   418 

   419 	Py_BEGIN_ALLOW_THREADS

   420 	result = CreateProcess(application_name,

   421 			       command_line,

   422 			       NULL,

   423 			       NULL,

   424 			       inherit_handles,

   425 			       creation_flags,

   426 			       environment ? PyString_AS_STRING(environment) : NULL,

   427 			       current_directory,

   428 			       &si,

   429 			       &pi);

   430 	Py_END_ALLOW_THREADS

   431 

   432 	Py_XDECREF(environment);

   433 

   434 	if (! result)

   435 		return PyErr_SetFromWindowsErr(GetLastError());

   436 

   437 	return Py_BuildValue("NNii",

   438 			     sp_handle_new(pi.hProcess),

   439 			     sp_handle_new(pi.hThread),

   440 			     pi.dwProcessId,

   441 			     pi.dwThreadId);

   442 }

   443 

   444 static PyObject *

   445 sp_TerminateProcess(PyObject* self, PyObject* args)

   446 {

   447 	BOOL result;

   448 

   449 	HANDLE process;

   450 	int exit_code;

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

   452 			       &process, &exit_code))

   453 		return NULL;

   454 

   455 	result = TerminateProcess(process, exit_code);

   456 

   457 	if (! result)

   458 		return PyErr_SetFromWindowsErr(GetLastError());

   459 

   460 	Py_INCREF(Py_None);

   461 	return Py_None;

   462 }

   463 

   464 static PyObject *

   465 sp_GetExitCodeProcess(PyObject* self, PyObject* args)

   466 {

   467 	DWORD exit_code;

   468 	BOOL result;

   469 

   470 	HANDLE process;

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

   472 		return NULL;

   473 

   474 	result = GetExitCodeProcess(process, &exit_code);

   475 

   476 	if (! result)

   477 		return PyErr_SetFromWindowsErr(GetLastError());

   478 

   479 	return PyInt_FromLong(exit_code);

   480 }

   481 

   482 static PyObject *

   483 sp_WaitForSingleObject(PyObject* self, PyObject* args)

   484 {

   485 	DWORD result;

   486 

   487 	HANDLE handle;

   488 	int milliseconds;

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

   490 	                  	     &handle,

   491 	                  	     &milliseconds))

   492 		return NULL;

   493 

   494 	Py_BEGIN_ALLOW_THREADS

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

   496 	Py_END_ALLOW_THREADS

   497 

   498 	if (result == WAIT_FAILED)

   499 		return PyErr_SetFromWindowsErr(GetLastError());

   500 

   501 	return PyInt_FromLong((int) result);

   502 }

   503 

   504 static PyObject *

   505 sp_GetVersion(PyObject* self, PyObject* args)

   506 {

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

   508 		return NULL;

   509 

   510 	return PyInt_FromLong((int) GetVersion());

   511 }

   512 

   513 static PyObject *

   514 sp_GetModuleFileName(PyObject* self, PyObject* args)

   515 {

   516 	BOOL result;

   517 	HMODULE module;

   518 	TCHAR filename[MAX_PATH];

   519 

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

   521 			       &module))

   522 		return NULL;

   523 

   524 	result = GetModuleFileName(module, filename, MAX_PATH);

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

   526 

   527 	if (! result)

   528 		return PyErr_SetFromWindowsErr(GetLastError());

   529 

   530 	return PyString_FromString(filename);

   531 }

   532 

   533 static PyMethodDef sp_functions[] = {

   534 	{"GetStdHandle",	sp_GetStdHandle,	METH_VARARGS},

   535 	{"GetCurrentProcess",	sp_GetCurrentProcess,	METH_VARARGS},

   536 	{"DuplicateHandle",	sp_DuplicateHandle,	METH_VARARGS},

   537 	{"CreatePipe",		sp_CreatePipe,		METH_VARARGS},

   538 	{"CreateProcess",	sp_CreateProcess,	METH_VARARGS},

   539 	{"TerminateProcess",	sp_TerminateProcess,	METH_VARARGS},

   540 	{"GetExitCodeProcess",	sp_GetExitCodeProcess,	METH_VARARGS},

   541 	{"WaitForSingleObject",	sp_WaitForSingleObject, METH_VARARGS},

   542 	{"GetVersion",		sp_GetVersion,		METH_VARARGS},

   543 	{"GetModuleFileName",	sp_GetModuleFileName,	METH_VARARGS},

   544 	{NULL, NULL}

   545 };

   546 

   547 /* -------------------------------------------------------------------- */

   548 

   549 static void

   550 defint(PyObject* d, const char* name, int value)

   551 {

   552 	PyObject* v = PyInt_FromLong((long) value);

   553 	if (v) {

   554 		PyDict_SetItemString(d, (char*) name, v);

   555 		Py_DECREF(v);

   556 	}

   557 }

   558 

   559 #if PY_VERSION_HEX >= 0x02030000

   560 PyMODINIT_FUNC

   561 #else

   562 DL_EXPORT(void)

   563 #endif

   564 init_subprocess()

   565 {

   566 	PyObject *d;

   567 	PyObject *m;

   568 

   569 	/* patch up object descriptors */

   570 	sp_handle_type.ob_type = &PyType_Type;

   571 	sp_handle_as_number.nb_int = (unaryfunc) sp_handle_as_int;

   572 

   573 	m = Py_InitModule("_subprocess", sp_functions);

   574 	if (m == NULL)

   575 		return;

   576 	d = PyModule_GetDict(m);

   577 

   578 	/* constants */

   579 	defint(d, "STD_INPUT_HANDLE", STD_INPUT_HANDLE);

   580 	defint(d, "STD_OUTPUT_HANDLE", STD_OUTPUT_HANDLE);

   581 	defint(d, "STD_ERROR_HANDLE", STD_ERROR_HANDLE);

   582 	defint(d, "DUPLICATE_SAME_ACCESS", DUPLICATE_SAME_ACCESS);

   583 	defint(d, "STARTF_USESTDHANDLES", STARTF_USESTDHANDLES);

   584 	defint(d, "STARTF_USESHOWWINDOW", STARTF_USESHOWWINDOW);

   585 	defint(d, "SW_HIDE", SW_HIDE);

   586 	defint(d, "INFINITE", INFINITE);

   587 	defint(d, "WAIT_OBJECT_0", WAIT_OBJECT_0);

   588 	defint(d, "CREATE_NEW_CONSOLE", CREATE_NEW_CONSOLE);

   589 }