PC/_subprocess.c

   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 = INVALID_HANDLE_VALUE;
  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 }