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         PyObject_FREE(self);
 108 }
 109
 110 static PyMethodDef sp_handle_methods[] = {
 111         {"Detach", (PyCFunction) sp_handle_detach, METH_VARARGS},
 112         {"Close",  (PyCFunction) sp_handle_close,  METH_VARARGS},
 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         int ret;
 254
 255         value = PyObject_GetAttrString(obj, name);
 256         if (! value) {
 257                 PyErr_Clear(); /* FIXME: propagate error? */
 258                 return 0;
 259         }
 260         ret = (int) PyInt_AsLong(value);
 261         Py_DECREF(value);
 262         return ret;
 263 }
 264
 265 static HANDLE
 266 gethandle(PyObject* obj, char* name)
 267 {
 268         sp_handle_object* value;
 269         HANDLE ret;
 270
 271         value = (sp_handle_object*) PyObject_GetAttrString(obj, name);
 272         if (! value) {
 273                 PyErr_Clear(); /* FIXME: propagate error? */
 274                 return NULL;
 275         }
 276         if (value->ob_type != &sp_handle_type)
 277                 ret = NULL;
 278         else
 279                 ret = value->handle;
 280         Py_DECREF(value);
 281         return ret;
 282 }
 283
 284 static PyObject*
 285 getenvironment(PyObject* environment)
 286 {
 287         int i, envsize;
 288         PyObject* out = NULL;
 289         PyObject* keys;
 290         PyObject* values;
 291         char* p;
 292
 293         /* convert environment dictionary to windows enviroment string */
 294         if (! PyMapping_Check(environment)) {
 295                 PyErr_SetString(
 296                     PyExc_TypeError, "environment must be dictionary or None");
 297                 return NULL;
 298         }
 299
 300         envsize = PyMapping_Length(environment);
 301
 302         keys = PyMapping_Keys(environment);
 303         values = PyMapping_Values(environment);
 304         if (!keys || !values)
 305                 goto error;
 306
 307         out = PyString_FromStringAndSize(NULL, 2048);
 308         if (! out)
 309                 goto error;
 310
 311         p = PyString_AS_STRING(out);
 312
 313         for (i = 0; i < envsize; i++) {
 314                 int ksize, vsize, totalsize;
 315                 PyObject* key = PyList_GET_ITEM(keys, i);
 316                 PyObject* value = PyList_GET_ITEM(values, i);
 317
 318                 if (! PyString_Check(key) || ! PyString_Check(value)) {
 319                         PyErr_SetString(PyExc_TypeError,
 320                                 "environment can only contain strings");
 321                         goto error;
 322                 }
 323                 ksize = PyString_GET_SIZE(key);
 324                 vsize = PyString_GET_SIZE(value);
 325                 totalsize = (p - PyString_AS_STRING(out)) + ksize + 1 +
 326                                                              vsize + 1 + 1;
 327                 if (totalsize > PyString_GET_SIZE(out)) {
 328                         int offset = p - PyString_AS_STRING(out);
 329                         _PyString_Resize(&out, totalsize + 1024);
 330                         p = PyString_AS_STRING(out) + offset;
 331                 }
 332                 memcpy(p, PyString_AS_STRING(key), ksize);
 333                 p += ksize;
 334                 *p++ = '=';
 335                 memcpy(p, PyString_AS_STRING(value), vsize);
 336                 p += vsize;
 337                 *p++ = '\0';
 338         }
 339
 340         /* add trailing null byte */
 341         *p++ = '\0';
 342         _PyString_Resize(&out, p - PyString_AS_STRING(out));
 343
 344         /* PyObject_Print(out, stdout, 0); */
 345
 346         Py_XDECREF(keys);
 347         Py_XDECREF(values);
 348
 349         return out;
 350
 351  error:
 352         Py_XDECREF(out);
 353         Py_XDECREF(keys);
 354         Py_XDECREF(values);
 355         return NULL;
 356 }
 357
 358 static PyObject *
 359 sp_CreateProcess(PyObject* self, PyObject* args)
 360 {
 361         BOOL result;
 362         PROCESS_INFORMATION pi;
 363         STARTUPINFO si;
 364         PyObject* environment;
 365
 366         char* application_name;
 367         char* command_line;
 368         PyObject* process_attributes; /* ignored */
 369         PyObject* thread_attributes; /* ignored */
 370         int inherit_handles;
 371         int creation_flags;
 372         PyObject* env_mapping;
 373         char* current_directory;
 374         PyObject* startup_info;
 375
 376         if (! PyArg_ParseTuple(args, "zzOOiiOzO:CreateProcess",
 377                                &application_name,
 378                                &command_line,
 379                                &process_attributes,
 380                                &thread_attributes,
 381                                &inherit_handles,
 382                                &creation_flags,
 383                                &env_mapping,
 384                                &current_directory,
 385                                &startup_info))
 386                 return NULL;
 387
 388         ZeroMemory(&si, sizeof(si));
 389         si.cb = sizeof(si);
 390
 391         /* note: we only support a small subset of all SI attributes */
 392         si.dwFlags = getint(startup_info, "dwFlags");
 393         si.wShowWindow = getint(startup_info, "wShowWindow");
 394         si.hStdInput = gethandle(startup_info, "hStdInput");
 395         si.hStdOutput = gethandle(startup_info, "hStdOutput");
 396         si.hStdError = gethandle(startup_info, "hStdError");
 397
 398         if (PyErr_Occurred())
 399                 return NULL;
 400
 401         if (env_mapping == Py_None)
 402                 environment = NULL;
 403         else {
 404                 environment = getenvironment(env_mapping);
 405                 if (! environment)
 406                         return NULL;
 407         }
 408
 409         Py_BEGIN_ALLOW_THREADS
 410         result = CreateProcess(application_name,
 411                                command_line,
 412                                NULL,
 413                                NULL,
 414                                inherit_handles,
 415                                creation_flags,
 416                                environment ? PyString_AS_STRING(environment) : NULL,
 417                                current_directory,
 418                                &si,
 419                                &pi);
 420         Py_END_ALLOW_THREADS
 421
 422         Py_XDECREF(environment);
 423
 424         if (! result)
 425                 return PyErr_SetFromWindowsErr(GetLastError());
 426
 427         return Py_BuildValue("NNii",
 428                              sp_handle_new(pi.hProcess),
 429                              sp_handle_new(pi.hThread),
 430                              pi.dwProcessId,
 431                              pi.dwThreadId);
 432 }
 433
 434 static PyObject *
 435 sp_TerminateProcess(PyObject* self, PyObject* args)
 436 {
 437         BOOL result;
 438
 439         long process;
 440         int exit_code;
 441         if (! PyArg_ParseTuple(args, "li:TerminateProcess", &process,
 442                                &exit_code))
 443                 return NULL;
 444
 445         result = TerminateProcess((HANDLE) process, exit_code);
 446
 447         if (! result)
 448                 return PyErr_SetFromWindowsErr(GetLastError());
 449
 450         Py_INCREF(Py_None);
 451         return Py_None;
 452 }
 453
 454 static PyObject *
 455 sp_GetExitCodeProcess(PyObject* self, PyObject* args)
 456 {
 457         DWORD exit_code;
 458         BOOL result;
 459
 460         long process;
 461         if (! PyArg_ParseTuple(args, "l:GetExitCodeProcess", &process))
 462                 return NULL;
 463
 464         result = GetExitCodeProcess((HANDLE) process, &exit_code);
 465
 466         if (! result)
 467                 return PyErr_SetFromWindowsErr(GetLastError());
 468
 469         return PyInt_FromLong(exit_code);
 470 }
 471
 472 static PyObject *
 473 sp_WaitForSingleObject(PyObject* self, PyObject* args)
 474 {
 475         DWORD result;
 476
 477         long handle;
 478         int milliseconds;
 479         if (! PyArg_ParseTuple(args, "li:WaitForSingleObject",
 480                                      &handle,
 481                                      &milliseconds))
 482                 return NULL;
 483
 484         Py_BEGIN_ALLOW_THREADS
 485         result = WaitForSingleObject((HANDLE) handle, (DWORD) milliseconds);
 486         Py_END_ALLOW_THREADS
 487
 488         if (result == WAIT_FAILED)
 489                 return PyErr_SetFromWindowsErr(GetLastError());
 490
 491         return PyInt_FromLong((int) result);
 492 }
 493
 494 static PyObject *
 495 sp_GetVersion(PyObject* self, PyObject* args)
 496 {
 497         if (! PyArg_ParseTuple(args, ":GetVersion"))
 498                 return NULL;
 499
 500         return PyInt_FromLong((int) GetVersion());
 501 }
 502
 503 static PyObject *
 504 sp_GetModuleFileName(PyObject* self, PyObject* args)
 505 {
 506         BOOL result;
 507         long module;
 508         TCHAR filename[MAX_PATH];
 509
 510         if (! PyArg_ParseTuple(args, "l:GetModuleFileName", &module))
 511                 return NULL;
 512
 513         result = GetModuleFileName((HMODULE)module, filename, MAX_PATH);
 514         filename[MAX_PATH-1] = '\0';
 515
 516         if (! result)
 517                 return PyErr_SetFromWindowsErr(GetLastError());
 518
 519         return PyString_FromString(filename);
 520 }
 521
 522 static PyMethodDef sp_functions[] = {
 523         {"GetStdHandle",        sp_GetStdHandle,        METH_VARARGS},
 524         {"GetCurrentProcess",   sp_GetCurrentProcess,   METH_VARARGS},
 525         {"DuplicateHandle",     sp_DuplicateHandle,     METH_VARARGS},
 526         {"CreatePipe",          sp_CreatePipe,          METH_VARARGS},
 527         {"CreateProcess",       sp_CreateProcess,       METH_VARARGS},
 528         {"TerminateProcess",    sp_TerminateProcess,    METH_VARARGS},
 529         {"GetExitCodeProcess",  sp_GetExitCodeProcess,  METH_VARARGS},
 530         {"WaitForSingleObject", sp_WaitForSingleObject, METH_VARARGS},
 531         {"GetVersion",          sp_GetVersion,          METH_VARARGS},
 532         {"GetModuleFileName",   sp_GetModuleFileName,   METH_VARARGS},
 533         {NULL, NULL}
 534 };
 535
 536 /* -------------------------------------------------------------------- */
 537
 538 static void
 539 defint(PyObject* d, const char* name, int value)
 540 {
 541         PyObject* v = PyInt_FromLong((long) value);
 542         if (v) {
 543                 PyDict_SetItemString(d, (char*) name, v);
 544                 Py_DECREF(v);
 545         }
 546 }
 547
 548 #if PY_VERSION_HEX >= 0x02030000
 549 PyMODINIT_FUNC
 550 #else
 551 DL_EXPORT(void)
 552 #endif
 553 init_subprocess()
 554 {
 555         PyObject *d;
 556         PyObject *m;
 557
 558         /* patch up object descriptors */
 559         sp_handle_type.ob_type = &PyType_Type;
 560         sp_handle_as_number.nb_int = (unaryfunc) sp_handle_as_int;
 561
 562         m = Py_InitModule("_subprocess", sp_functions);
 563         if (m == NULL)
 564                 return;
 565         d = PyModule_GetDict(m);
 566
 567         /* constants */
 568         defint(d, "STD_INPUT_HANDLE", STD_INPUT_HANDLE);
 569         defint(d, "STD_OUTPUT_HANDLE", STD_OUTPUT_HANDLE);
 570         defint(d, "STD_ERROR_HANDLE", STD_ERROR_HANDLE);
 571         defint(d, "DUPLICATE_SAME_ACCESS", DUPLICATE_SAME_ACCESS);
 572         defint(d, "STARTF_USESTDHANDLES", STARTF_USESTDHANDLES);
 573         defint(d, "STARTF_USESHOWWINDOW", STARTF_USESHOWWINDOW);
 574         defint(d, "SW_HIDE", SW_HIDE);
 575         defint(d, "INFINITE", INFINITE);
 576         defint(d, "WAIT_OBJECT_0", WAIT_OBJECT_0);
 577         defint(d, "CREATE_NEW_CONSOLE", CREATE_NEW_CONSOLE);
 578 }