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 static PyObject*
  73 sp_handle_detach(sp_handle_object* self, PyObject* args)
  74 {
  75         HANDLE handle;
  76
  77         if (! PyArg_ParseTuple(args, ":Detach"))
  78                 return NULL;
  79
  80         handle = self->handle;
  81
  82         self->handle = NULL;
  83
  84         /* note: return the current handle, as an integer */
  85         return PyInt_FromLong((long) handle);
  86 }
  87
  88 static PyObject*
  89 sp_handle_close(sp_handle_object* self, PyObject* args)
  90 {
  91         if (! PyArg_ParseTuple(args, ":Close"))
  92                 return NULL;
  93
  94         if (self->handle != INVALID_HANDLE_VALUE) {
  95                 CloseHandle(self->handle);
  96                 self->handle = INVALID_HANDLE_VALUE;
  97         }
  98         Py_INCREF(Py_None);
  99         return Py_None;
 100 }
 101
 102 static void
 103 sp_handle_dealloc(sp_handle_object* self)
 104 {
 105         if (self->handle != INVALID_HANDLE_VALUE)
 106                 CloseHandle(self->handle);
 107         PyMem_DEL(self);
 108 }
 109
 110 static PyMethodDef sp_handle_methods[] = {
 111         {"Detach", (PyCFunction) sp_handle_detach, 1},
 112         {"Close", (PyCFunction) sp_handle_close, 1},
 113         {NULL, NULL}
 114 };
 115
 116 static PyObject*
 117 sp_handle_getattr(sp_handle_object* self, char* name)
 118 {
 119         return Py_FindMethod(sp_handle_methods, (PyObject*) self, name);
 120 }
 121
 122 static PyObject*
 123 sp_handle_as_int(sp_handle_object* self)
 124 {
 125         return PyInt_FromLong((long) self->handle);
 126 }
 127
 128 static PyNumberMethods sp_handle_as_number;
 129
 130 statichere PyTypeObject sp_handle_type = {
 131         PyObject_HEAD_INIT(NULL)
 132         0,                              /*ob_size*/
 133         "_subprocess_handle", sizeof(sp_handle_object), 0,
 134         (destructor) sp_handle_dealloc, /*tp_dealloc*/
 135         0, /*tp_print*/
 136         (getattrfunc) sp_handle_getattr,/*tp_getattr*/
 137         0,                              /*tp_setattr*/
 138         0,                              /*tp_compare*/
 139         0,                              /*tp_repr*/
 140         &sp_handle_as_number,           /*tp_as_number */
 141         0,                              /*tp_as_sequence */
 142         0,                              /*tp_as_mapping */
 143         0                               /*tp_hash*/
 144 };
 145
 146 /* -------------------------------------------------------------------- */
 147 /* windows API functions */
 148
 149 static PyObject *
 150 sp_GetStdHandle(PyObject* self, PyObject* args)
 151 {
 152         HANDLE handle;
 153         int std_handle;
 154
 155         if (! PyArg_ParseTuple(args, "i:GetStdHandle", &std_handle))
 156                 return NULL;
 157
 158         Py_BEGIN_ALLOW_THREADS
 159         handle = GetStdHandle((DWORD) std_handle);
 160         Py_END_ALLOW_THREADS
 161
 162         if (handle == INVALID_HANDLE_VALUE)
 163                 return PyErr_SetFromWindowsErr(GetLastError());
 164
 165         if (! handle) {
 166                 Py_INCREF(Py_None);
 167                 return Py_None;
 168         }
 169
 170         /* note: returns integer, not handle object */
 171         return PyInt_FromLong((long) handle);
 172 }
 173
 174 static PyObject *
 175 sp_GetCurrentProcess(PyObject* self, PyObject* args)
 176 {
 177         if (! PyArg_ParseTuple(args, ":GetCurrentProcess"))
 178                 return NULL;
 179
 180         return sp_handle_new(GetCurrentProcess());
 181 }
 182
 183 static PyObject *
 184 sp_DuplicateHandle(PyObject* self, PyObject* args)
 185 {
 186         HANDLE target_handle;
 187         BOOL result;
 188
 189         long source_process_handle;
 190         long source_handle;
 191         long target_process_handle;
 192         int desired_access;
 193         int inherit_handle;
 194         int options = 0;
 195
 196         if (! PyArg_ParseTuple(args, "lllii|i:DuplicateHandle",
 197                                &source_process_handle,
 198                                &source_handle,
 199                                &target_process_handle,
 200                                &desired_access,
 201                                &inherit_handle,
 202                                &options))
 203                 return NULL;
 204
 205         Py_BEGIN_ALLOW_THREADS
 206         result = DuplicateHandle(
 207                 (HANDLE) source_process_handle,
 208                 (HANDLE) source_handle,
 209                 (HANDLE) target_process_handle,
 210                 &target_handle,
 211                 desired_access,
 212                 inherit_handle,
 213                 options
 214         );
 215         Py_END_ALLOW_THREADS
 216
 217         if (! result)
 218                 return PyErr_SetFromWindowsErr(GetLastError());
 219
 220         return sp_handle_new(target_handle);
 221 }
 222
 223 static PyObject *
 224 sp_CreatePipe(PyObject* self, PyObject* args)
 225 {
 226         HANDLE read_pipe;
 227         HANDLE write_pipe;
 228         BOOL result;
 229
 230         PyObject* pipe_attributes; /* ignored */
 231         int size;
 232
 233         if (! PyArg_ParseTuple(args, "Oi:CreatePipe", &pipe_attributes, &size))
 234                 return NULL;
 235
 236         Py_BEGIN_ALLOW_THREADS
 237         result = CreatePipe(&read_pipe, &write_pipe, NULL, size);
 238         Py_END_ALLOW_THREADS
 239
 240         if (! result)
 241                 return PyErr_SetFromWindowsErr(GetLastError());
 242
 243         return Py_BuildValue(
 244                 "NN", sp_handle_new(read_pipe), sp_handle_new(write_pipe));
 245 }
 246
 247 /* helpers for createprocess */
 248
 249 static int
 250 getint(PyObject* obj, char* name)
 251 {
 252         PyObject* value;
 253
 254         value = PyObject_GetAttrString(obj, name);
 255         if (! value) {
 256                 PyErr_Clear(); /* FIXME: propagate error? */
 257                 return 0;
 258         }
 259         return (int) PyInt_AsLong(value);
 260 }
 261
 262 static HANDLE
 263 gethandle(PyObject* obj, char* name)
 264 {
 265         sp_handle_object* value;
 266
 267         value = (sp_handle_object*) PyObject_GetAttrString(obj, name);
 268         if (! value) {
 269                 PyErr_Clear(); /* FIXME: propagate error? */
 270                 return NULL;
 271         }
 272         if (value->ob_type != &sp_handle_type)
 273                 return NULL;
 274         return value->handle;
 275 }
 276
 277 static PyObject*
 278 getenvironment(PyObject* environment)
 279 {
 280         int i, envsize;
 281         PyObject* out = NULL;
 282         PyObject* keys;
 283         PyObject* values;
 284         char* p;
 285
 286         /* convert environment dictionary to windows enviroment string */
 287         if (! PyMapping_Check(environment)) {
 288                 PyErr_SetString(
 289                     PyExc_TypeError, "environment must be dictionary or None");
 290                 return NULL;
 291         }
 292
 293         envsize = PyMapping_Length(environment);
 294
 295         keys = PyMapping_Keys(environment);
 296         values = PyMapping_Values(environment);
 297         if (!keys || !values)
 298                 goto error;
 299
 300         out = PyString_FromStringAndSize(NULL, 2048);
 301         if (! out)
 302                 goto error;
 303
 304         p = PyString_AS_STRING(out);
 305
 306         for (i = 0; i < envsize; i++) {
 307                 int ksize, vsize, totalsize;
 308                 PyObject* key = PyList_GET_ITEM(keys, i);
 309                 PyObject* value = PyList_GET_ITEM(values, i);
 310
 311                 if (! PyString_Check(key) || ! PyString_Check(value)) {
 312                         PyErr_SetString(PyExc_TypeError,
 313                                 "environment can only contain strings");
 314                         goto error;
 315                 }
 316                 ksize = PyString_GET_SIZE(key);
 317                 vsize = PyString_GET_SIZE(value);
 318                 totalsize = (p - PyString_AS_STRING(out)) + ksize + 1 +
 319                                                              vsize + 1 + 1;
 320                 if (totalsize > PyString_GET_SIZE(out)) {
 321                         int offset = p - PyString_AS_STRING(out);
 322                         _PyString_Resize(&out, totalsize + 1024);
 323                         p = PyString_AS_STRING(out) + offset;
 324                 }
 325                 memcpy(p, PyString_AS_STRING(key), ksize);
 326                 p += ksize;
 327                 *p++ = '=';
 328                 memcpy(p, PyString_AS_STRING(value), vsize);
 329                 p += vsize;
 330                 *p++ = '\0';
 331         }
 332
 333         /* add trailing null byte */
 334         *p++ = '\0';
 335         _PyString_Resize(&out, p - PyString_AS_STRING(out));
 336
 337         /* PyObject_Print(out, stdout, 0); */
 338
 339         Py_XDECREF(keys);
 340         Py_XDECREF(values);
 341
 342         return out;
 343
 344  error:
 345         Py_XDECREF(out);
 346         Py_XDECREF(keys);
 347         Py_XDECREF(values);
 348         return NULL;
 349 }
 350
 351 static PyObject *
 352 sp_CreateProcess(PyObject* self, PyObject* args)
 353 {
 354         BOOL result;
 355         PROCESS_INFORMATION pi;
 356         STARTUPINFO si;
 357         PyObject* environment;
 358
 359         char* application_name;
 360         char* command_line;
 361         PyObject* process_attributes; /* ignored */
 362         PyObject* thread_attributes; /* ignored */
 363         int inherit_handles;
 364         int creation_flags;
 365         PyObject* env_mapping;
 366         char* current_directory;
 367         PyObject* startup_info;
 368
 369         if (! PyArg_ParseTuple(args, "zzOOiiOzO:CreateProcess",
 370                                &application_name,
 371                                &command_line,
 372                                &process_attributes,
 373                                &thread_attributes,
 374                                &inherit_handles,
 375                                &creation_flags,
 376                                &env_mapping,
 377                                &current_directory,
 378                                &startup_info))
 379                 return NULL;
 380
 381         ZeroMemory(&si, sizeof(si));
 382         si.cb = sizeof(si);
 383
 384         /* note: we only support a small subset of all SI attributes */
 385         si.dwFlags = getint(startup_info, "dwFlags");
 386         si.wShowWindow = getint(startup_info, "wShowWindow");
 387         si.hStdInput = gethandle(startup_info, "hStdInput");
 388         si.hStdOutput = gethandle(startup_info, "hStdOutput");
 389         si.hStdError = gethandle(startup_info, "hStdError");
 390
 391         if (PyErr_Occurred())
 392                 return NULL;
 393
 394         if (env_mapping == Py_None)
 395                 environment = NULL;
 396         else {
 397                 environment = getenvironment(env_mapping);
 398                 if (! environment)
 399                         return NULL;
 400         }
 401
 402         Py_BEGIN_ALLOW_THREADS
 403         result = CreateProcess(application_name,
 404                                command_line,
 405                                NULL,
 406                                NULL,
 407                                inherit_handles,
 408                                creation_flags,
 409                                environment ? PyString_AS_STRING(environment) : NULL,
 410                                current_directory,
 411                                &si,
 412                                &pi);
 413         Py_END_ALLOW_THREADS
 414
 415         Py_XDECREF(environment);
 416
 417         if (! result)
 418                 return PyErr_SetFromWindowsErr(GetLastError());
 419
 420         return Py_BuildValue("NNii",
 421                              sp_handle_new(pi.hProcess),
 422                              sp_handle_new(pi.hThread),
 423                              pi.dwProcessId,
 424                              pi.dwThreadId);
 425 }
 426
 427 static PyObject *
 428 sp_TerminateProcess(PyObject* self, PyObject* args)
 429 {
 430         BOOL result;
 431
 432         long process;
 433         int exit_code;
 434         if (! PyArg_ParseTuple(args, "li:TerminateProcess", &process,
 435                                &exit_code))
 436                 return NULL;
 437
 438         result = TerminateProcess((HANDLE) process, exit_code);
 439
 440         if (! result)
 441                 return PyErr_SetFromWindowsErr(GetLastError());
 442
 443         Py_INCREF(Py_None);
 444         return Py_None;
 445 }
 446
 447 static PyObject *
 448 sp_GetExitCodeProcess(PyObject* self, PyObject* args)
 449 {
 450         DWORD exit_code;
 451         BOOL result;
 452
 453         long process;
 454         if (! PyArg_ParseTuple(args, "l:GetExitCodeProcess", &process))
 455                 return NULL;
 456
 457         result = GetExitCodeProcess((HANDLE) process, &exit_code);
 458
 459         if (! result)
 460                 return PyErr_SetFromWindowsErr(GetLastError());
 461
 462         return PyInt_FromLong(exit_code);
 463 }
 464
 465 static PyObject *
 466 sp_WaitForSingleObject(PyObject* self, PyObject* args)
 467 {
 468         DWORD result;
 469
 470         long handle;
 471         int milliseconds;
 472         if (! PyArg_ParseTuple(args, "li:WaitForSingleObject",
 473                                      &handle,
 474                                      &milliseconds))
 475                 return NULL;
 476
 477         Py_BEGIN_ALLOW_THREADS
 478         result = WaitForSingleObject((HANDLE) handle, (DWORD) milliseconds);
 479         Py_END_ALLOW_THREADS
 480
 481         if (result == WAIT_FAILED)
 482                 return PyErr_SetFromWindowsErr(GetLastError());
 483
 484         return PyInt_FromLong((int) result);
 485 }
 486
 487 static PyObject *
 488 sp_GetVersion(PyObject* self, PyObject* args)
 489 {
 490         if (! PyArg_ParseTuple(args, ":GetVersion"))
 491                 return NULL;
 492
 493         return PyInt_FromLong((int) GetVersion());
 494 }
 495
 496 static PyObject *
 497 sp_GetModuleFileName(PyObject* self, PyObject* args)
 498 {
 499         BOOL result;
 500         long module;
 501         TCHAR filename[MAX_PATH];
 502
 503         if (! PyArg_ParseTuple(args, "l:GetModuleFileName", &module))
 504                 return NULL;
 505
 506         result = GetModuleFileName((HMODULE)module, filename, MAX_PATH);
 507         filename[MAX_PATH-1] = '\0';
 508
 509         if (! result)
 510                 return PyErr_SetFromWindowsErr(GetLastError());
 511
 512         return PyString_FromString(filename);
 513 }
 514
 515 static PyMethodDef sp_functions[] = {
 516         {"GetStdHandle",        sp_GetStdHandle,        METH_VARARGS},
 517         {"GetCurrentProcess",   sp_GetCurrentProcess,   METH_VARARGS},
 518         {"DuplicateHandle",     sp_DuplicateHandle,     METH_VARARGS},
 519         {"CreatePipe",          sp_CreatePipe,          METH_VARARGS},
 520         {"CreateProcess",       sp_CreateProcess,       METH_VARARGS},
 521         {"TerminateProcess",    sp_TerminateProcess,    METH_VARARGS},
 522         {"GetExitCodeProcess",  sp_GetExitCodeProcess,  METH_VARARGS},
 523         {"WaitForSingleObject", sp_WaitForSingleObject, METH_VARARGS},
 524         {"GetVersion",          sp_GetVersion,          METH_VARARGS},
 525         {"GetModuleFileName",   sp_GetModuleFileName,   METH_VARARGS},
 526         {NULL, NULL}
 527 };
 528
 529 /* -------------------------------------------------------------------- */
 530
 531 static void
 532 defint(PyObject* d, const char* name, int value)
 533 {
 534         PyObject* v = PyInt_FromLong((long) value);
 535         if (v) {
 536                 PyDict_SetItemString(d, (char*) name, v);
 537                 Py_DECREF(v);
 538         }
 539 }
 540
 541 #if PY_VERSION_HEX >= 0x02030000
 542 PyMODINIT_FUNC
 543 #else
 544 DL_EXPORT(void)
 545 #endif
 546 init_subprocess()
 547 {
 548         PyObject *d;
 549         PyObject *m;
 550
 551         /* patch up object descriptors */
 552         sp_handle_type.ob_type = &PyType_Type;
 553         sp_handle_as_number.nb_int = (unaryfunc) sp_handle_as_int;
 554
 555         m = Py_InitModule("_subprocess", sp_functions);
 556         if (m == NULL)
 557                 return;
 558         d = PyModule_GetDict(m);
 559
 560         /* constants */
 561         defint(d, "STD_INPUT_HANDLE", STD_INPUT_HANDLE);
 562         defint(d, "STD_OUTPUT_HANDLE", STD_OUTPUT_HANDLE);
 563         defint(d, "STD_ERROR_HANDLE", STD_ERROR_HANDLE);
 564         defint(d, "DUPLICATE_SAME_ACCESS", DUPLICATE_SAME_ACCESS);
 565         defint(d, "STARTF_USESTDHANDLES", STARTF_USESTDHANDLES);
 566         defint(d, "STARTF_USESHOWWINDOW", STARTF_USESHOWWINDOW);
 567         defint(d, "SW_HIDE", SW_HIDE);
 568         defint(d, "INFINITE", INFINITE);
 569         defint(d, "WAIT_OBJECT_0", WAIT_OBJECT_0);
 570         defint(d, "CREATE_NEW_CONSOLE", CREATE_NEW_CONSOLE);
 571 }