1 """Windows-specific implementation of process utilities with direct WinAPI. 2 3 This file is meant to be used by process.py 4 """ 5 6 #----------------------------------------------------------------------------- 7 # Copyright (C) 2010-2011 The IPython Development Team 8 # 9 # Distributed under the terms of the BSD License. The full license is in 10 # the file COPYING, distributed as part of this software. 11 #----------------------------------------------------------------------------- 12 13 14 # stdlib 15 import os, sys, threading 16 import ctypes, msvcrt 17 18 # Win32 API types needed for the API calls 19 from ctypes import POINTER 20 from ctypes.wintypes import HANDLE, HLOCAL, LPVOID, WORD, DWORD, BOOL, \ 21 ULONG, LPCWSTR 22 LPDWORD = POINTER(DWORD) 23 LPHANDLE = POINTER(HANDLE) 24 ULONG_PTR = POINTER(ULONG) 25 class SECURITY_ATTRIBUTES(ctypes.Structure): 26 _fields_ = [("nLength", DWORD), 27 ("lpSecurityDescriptor", LPVOID), 28 ("bInheritHandle", BOOL)] 29 LPSECURITY_ATTRIBUTES = POINTER(SECURITY_ATTRIBUTES) 30 class STARTUPINFO(ctypes.Structure): 31 _fields_ = [("cb", DWORD), 32 ("lpReserved", LPCWSTR), 33 ("lpDesktop", LPCWSTR), 34 ("lpTitle", LPCWSTR), 35 ("dwX", DWORD), 36 ("dwY", DWORD), 37 ("dwXSize", DWORD), 38 ("dwYSize", DWORD), 39 ("dwXCountChars", DWORD), 40 ("dwYCountChars", DWORD), 41 ("dwFillAttribute", DWORD), 42 ("dwFlags", DWORD), 43 ("wShowWindow", WORD), 44 ("cbReserved2", WORD), 45 ("lpReserved2", LPVOID), 46 ("hStdInput", HANDLE), 47 ("hStdOutput", HANDLE), 48 ("hStdError", HANDLE)] 49 LPSTARTUPINFO = POINTER(STARTUPINFO) 50 class PROCESS_INFORMATION(ctypes.Structure): 51 _fields_ = [("hProcess", HANDLE), 52 ("hThread", HANDLE), 53 ("dwProcessId", DWORD), 54 ("dwThreadId", DWORD)] 55 LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION) 56 57 # Win32 API constants needed 58 ERROR_HANDLE_EOF = 38 59 ERROR_BROKEN_PIPE = 109 60 ERROR_NO_DATA = 232 61 HANDLE_FLAG_INHERIT = 0x0001 62 STARTF_USESTDHANDLES = 0x0100 63 CREATE_SUSPENDED = 0x0004 64 CREATE_NEW_CONSOLE = 0x0010 65 CREATE_NO_WINDOW = 0x08000000 66 STILL_ACTIVE = 259 67 WAIT_TIMEOUT = 0x0102 68 WAIT_FAILED = 0xFFFFFFFF 69 INFINITE = 0xFFFFFFFF 70 DUPLICATE_SAME_ACCESS = 0x00000002 71 ENABLE_ECHO_INPUT = 0x0004 72 ENABLE_LINE_INPUT = 0x0002 73 ENABLE_PROCESSED_INPUT = 0x0001 74 75 # Win32 API functions needed 76 GetLastError = ctypes.windll.kernel32.GetLastError 77 GetLastError.argtypes = [] 78 GetLastError.restype = DWORD 79 80 CreateFile = ctypes.windll.kernel32.CreateFileW 81 CreateFile.argtypes = [LPCWSTR, DWORD, DWORD, LPVOID, DWORD, DWORD, HANDLE] 82 CreateFile.restype = HANDLE 83 84 CreatePipe = ctypes.windll.kernel32.CreatePipe 85 CreatePipe.argtypes = [POINTER(HANDLE), POINTER(HANDLE), 86 LPSECURITY_ATTRIBUTES, DWORD] 87 CreatePipe.restype = BOOL 88 89 CreateProcess = ctypes.windll.kernel32.CreateProcessW 90 CreateProcess.argtypes = [LPCWSTR, LPCWSTR, LPSECURITY_ATTRIBUTES, 91 LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFO, 92 LPPROCESS_INFORMATION] 93 CreateProcess.restype = BOOL 94 95 GetExitCodeProcess = ctypes.windll.kernel32.GetExitCodeProcess 96 GetExitCodeProcess.argtypes = [HANDLE, LPDWORD] 97 GetExitCodeProcess.restype = BOOL 98 99 GetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess 100 GetCurrentProcess.argtypes = [] 101 GetCurrentProcess.restype = HANDLE 102 103 ResumeThread = ctypes.windll.kernel32.ResumeThread 104 ResumeThread.argtypes = [HANDLE] 105 ResumeThread.restype = DWORD 106 107 ReadFile = ctypes.windll.kernel32.ReadFile 108 ReadFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID] 109 ReadFile.restype = BOOL 110 111 WriteFile = ctypes.windll.kernel32.WriteFile 112 WriteFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID] 113 WriteFile.restype = BOOL 114 115 GetConsoleMode = ctypes.windll.kernel32.GetConsoleMode 116 GetConsoleMode.argtypes = [HANDLE, LPDWORD] 117 GetConsoleMode.restype = BOOL 118 119 SetConsoleMode = ctypes.windll.kernel32.SetConsoleMode 120 SetConsoleMode.argtypes = [HANDLE, DWORD] 121 SetConsoleMode.restype = BOOL 122 123 FlushConsoleInputBuffer = ctypes.windll.kernel32.FlushConsoleInputBuffer 124 FlushConsoleInputBuffer.argtypes = [HANDLE] 125 FlushConsoleInputBuffer.restype = BOOL 126 127 WaitForSingleObject = ctypes.windll.kernel32.WaitForSingleObject 128 WaitForSingleObject.argtypes = [HANDLE, DWORD] 129 WaitForSingleObject.restype = DWORD 130 131 DuplicateHandle = ctypes.windll.kernel32.DuplicateHandle 132 DuplicateHandle.argtypes = [HANDLE, HANDLE, HANDLE, LPHANDLE, 133 DWORD, BOOL, DWORD] 134 DuplicateHandle.restype = BOOL 135 136 SetHandleInformation = ctypes.windll.kernel32.SetHandleInformation 137 SetHandleInformation.argtypes = [HANDLE, DWORD, DWORD] 138 SetHandleInformation.restype = BOOL 139 140 CloseHandle = ctypes.windll.kernel32.CloseHandle 141 CloseHandle.argtypes = [HANDLE] 142 CloseHandle.restype = BOOL 143 144 CommandLineToArgvW = ctypes.windll.shell32.CommandLineToArgvW 145 CommandLineToArgvW.argtypes = [LPCWSTR, POINTER(ctypes.c_int)] 146 CommandLineToArgvW.restype = POINTER(LPCWSTR) 147 148 LocalFree = ctypes.windll.kernel32.LocalFree 149 LocalFree.argtypes = [HLOCAL] 150 LocalFree.restype = HLOCAL 151 152 class AvoidUNCPath(object): 153 """A context manager to protect command execution from UNC paths. 154 155 In the Win32 API, commands can't be invoked with the cwd being a UNC path. 156 This context manager temporarily changes directory to the 'C:' drive on 157 entering, and restores the original working directory on exit. 158 159 The context manager returns the starting working directory *if* it made a 160 change and None otherwise, so that users can apply the necessary adjustment 161 to their system calls in the event of a change. 162 163 Examples 164 -------- 165 :: 166 cmd = 'dir' 167 with AvoidUNCPath() as path: 168 if path is not None: 169 cmd = '"pushd %s &&"%s' % (path, cmd) 170 os.system(cmd) 171 """ 172 def __enter__(self): 173 self.path = os.getcwd() 174 self.is_unc_path = self.path.startswith(r"\\") 175 if self.is_unc_path: 176 # change to c drive (as cmd.exe cannot handle UNC addresses) 177 os.chdir("C:") 178 return self.path 179 else: 180 # We return None to signal that there was no change in the working 181 # directory 182 return None 183 184 def __exit__(self, exc_type, exc_value, traceback): 185 if self.is_unc_path: 186 os.chdir(self.path) 187 188 189 class Win32ShellCommandController(object): 190 """Runs a shell command in a 'with' context. 191 192 This implementation is Win32-specific. 193 194 Example: 195 # Runs the command interactively with default console stdin/stdout 196 with ShellCommandController('python -i') as scc: 197 scc.run() 198 199 # Runs the command using the provided functions for stdin/stdout 200 def my_stdout_func(s): 201 # print or save the string 's' 202 write_to_stdout(s) 203 def my_stdin_func(): 204 # If input is available, return it as a string. 205 if input_available(): 206 return get_input() 207 # If no input available, return None after a short delay to 208 # keep from blocking. 209 else: 210 time.sleep(0.01) 211 return None 212 213 with ShellCommandController('python -i') as scc: 214 scc.run(my_stdout_func, my_stdin_func) 215 """ 216 217 def __init__(self, cmd, mergeout = True): 218 """Initializes the shell command controller. 219 220 The cmd is the program to execute, and mergeout is 221 whether to blend stdout and stderr into one output 222 in stdout. Merging them together in this fashion more 223 reliably keeps stdout and stderr in the correct order 224 especially for interactive shell usage. 225 """ 226 self.cmd = cmd 227 self.mergeout = mergeout 228 229 def __enter__(self): 230 cmd = self.cmd 231 mergeout = self.mergeout 232 233 self.hstdout, self.hstdin, self.hstderr = None, None, None 234 self.piProcInfo = None 235 try: 236 p_hstdout, c_hstdout, p_hstderr, \ 237 c_hstderr, p_hstdin, c_hstdin = [None]*6 238 239 # SECURITY_ATTRIBUTES with inherit handle set to True 240 saAttr = SECURITY_ATTRIBUTES() 241 saAttr.nLength = ctypes.sizeof(saAttr) 242 saAttr.bInheritHandle = True 243 saAttr.lpSecurityDescriptor = None 244 245 def create_pipe(uninherit): 246 """Creates a Windows pipe, which consists of two handles. 247 248 The 'uninherit' parameter controls which handle is not 249 inherited by the child process. 250 """ 251 handles = HANDLE(), HANDLE() 252 if not CreatePipe(ctypes.byref(handles[0]), 253 ctypes.byref(handles[1]), ctypes.byref(saAttr), 0): 254 raise ctypes.WinError() 255 if not SetHandleInformation(handles[uninherit], 256 HANDLE_FLAG_INHERIT, 0): 257 raise ctypes.WinError() 258 return handles[0].value, handles[1].value 259 260 p_hstdout, c_hstdout = create_pipe(uninherit=0) 261 # 'mergeout' signals that stdout and stderr should be merged. 262 # We do that by using one pipe for both of them. 263 if mergeout: 264 c_hstderr = HANDLE() 265 if not DuplicateHandle(GetCurrentProcess(), c_hstdout, 266 GetCurrentProcess(), ctypes.byref(c_hstderr), 267 0, True, DUPLICATE_SAME_ACCESS): 268 raise ctypes.WinError() 269 else: 270 p_hstderr, c_hstderr = create_pipe(uninherit=0) 271 c_hstdin, p_hstdin = create_pipe(uninherit=1) 272 273 # Create the process object 274 piProcInfo = PROCESS_INFORMATION() 275 siStartInfo = STARTUPINFO() 276 siStartInfo.cb = ctypes.sizeof(siStartInfo) 277 siStartInfo.hStdInput = c_hstdin 278 siStartInfo.hStdOutput = c_hstdout 279 siStartInfo.hStdError = c_hstderr 280 siStartInfo.dwFlags = STARTF_USESTDHANDLES 281 dwCreationFlags = CREATE_SUSPENDED | CREATE_NO_WINDOW # | CREATE_NEW_CONSOLE 282 283 if not CreateProcess(None, 284 u"cmd.exe /c " + cmd, 285 None, None, True, dwCreationFlags, 286 None, None, ctypes.byref(siStartInfo), 287 ctypes.byref(piProcInfo)): 288 raise ctypes.WinError() 289 290 # Close this process's versions of the child handles 291 CloseHandle(c_hstdin) 292 c_hstdin = None 293 CloseHandle(c_hstdout) 294 c_hstdout = None 295 if c_hstderr is not None: 296 CloseHandle(c_hstderr) 297 c_hstderr = None 298 299 # Transfer ownership of the parent handles to the object 300 self.hstdin = p_hstdin 301 p_hstdin = None 302 self.hstdout = p_hstdout 303 p_hstdout = None 304 if not mergeout: 305 self.hstderr = p_hstderr 306 p_hstderr = None 307 self.piProcInfo = piProcInfo 308 309 finally: 310 if p_hstdin: 311 CloseHandle(p_hstdin) 312 if c_hstdin: 313 CloseHandle(c_hstdin) 314 if p_hstdout: 315 CloseHandle(p_hstdout) 316 if c_hstdout: 317 CloseHandle(c_hstdout) 318 if p_hstderr: 319 CloseHandle(p_hstderr) 320 if c_hstderr: 321 CloseHandle(c_hstderr) 322 323 return self 324 325 def _stdin_thread(self, handle, hprocess, func, stdout_func): 326 exitCode = DWORD() 327 bytesWritten = DWORD(0) 328 while True: 329 #print("stdin thread loop start") 330 # Get the input string (may be bytes or unicode) 331 data = func() 332 333 # None signals to poll whether the process has exited 334 if data is None: 335 #print("checking for process completion") 336 if not GetExitCodeProcess(hprocess, ctypes.byref(exitCode)): 337 raise ctypes.WinError() 338 if exitCode.value != STILL_ACTIVE: 339 return 340 # TESTING: Does zero-sized writefile help? 341 if not WriteFile(handle, "", 0, 342 ctypes.byref(bytesWritten), None): 343 raise ctypes.WinError() 344 continue 345 #print("\nGot str %s\n" % repr(data), file=sys.stderr) 346 347 # Encode the string to the console encoding 348 if isinstance(data, unicode): #FIXME: Python3 349 data = data.encode('utf_8') 350 351 # What we have now must be a string of bytes 352 if not isinstance(data, str): #FIXME: Python3 353 raise RuntimeError("internal stdin function string error") 354 355 # An empty string signals EOF 356 if len(data) == 0: 357 return 358 359 # In a windows console, sometimes the input is echoed, 360 # but sometimes not. How do we determine when to do this? 361 stdout_func(data) 362 # WriteFile may not accept all the data at once. 363 # Loop until everything is processed 364 while len(data) != 0: 365 #print("Calling writefile") 366 if not WriteFile(handle, data, len(data), 367 ctypes.byref(bytesWritten), None): 368 # This occurs at exit 369 if GetLastError() == ERROR_NO_DATA: 370 return 371 raise ctypes.WinError() 372 #print("Called writefile") 373 data = data[bytesWritten.value:] 374 375 def _stdout_thread(self, handle, func): 376 # Allocate the output buffer 377 data = ctypes.create_string_buffer(4096) 378 while True: 379 bytesRead = DWORD(0) 380 if not ReadFile(handle, data, 4096, 381 ctypes.byref(bytesRead), None): 382 le = GetLastError() 383 if le == ERROR_BROKEN_PIPE: 384 return 385 else: 386 raise ctypes.WinError() 387 # FIXME: Python3 388 s = data.value[0:bytesRead.value] 389 #print("\nv: %s" % repr(s), file=sys.stderr) 390 func(s.decode('utf_8', 'replace')) 391 392 def run(self, stdout_func = None, stdin_func = None, stderr_func = None): 393 """Runs the process, using the provided functions for I/O. 394 395 The function stdin_func should return strings whenever a 396 character or characters become available. 397 The functions stdout_func and stderr_func are called whenever 398 something is printed to stdout or stderr, respectively. 399 These functions are called from different threads (but not 400 concurrently, because of the GIL). 401 """ 402 if stdout_func is None and stdin_func is None and stderr_func is None: 403 return self._run_stdio() 404 405 if stderr_func is not None and self.mergeout: 406 raise RuntimeError("Shell command was initiated with " 407 "merged stdin/stdout, but a separate stderr_func " 408 "was provided to the run() method") 409 410 # Create a thread for each input/output handle 411 stdin_thread = None 412 threads = [] 413 if stdin_func: 414 stdin_thread = threading.Thread(target=self._stdin_thread, 415 args=(self.hstdin, self.piProcInfo.hProcess, 416 stdin_func, stdout_func)) 417 threads.append(threading.Thread(target=self._stdout_thread, 418 args=(self.hstdout, stdout_func))) 419 if not self.mergeout: 420 if stderr_func is None: 421 stderr_func = stdout_func 422 threads.append(threading.Thread(target=self._stdout_thread, 423 args=(self.hstderr, stderr_func))) 424 # Start the I/O threads and the process 425 if ResumeThread(self.piProcInfo.hThread) == 0xFFFFFFFF: 426 raise ctypes.WinError() 427 if stdin_thread is not None: 428 stdin_thread.start() 429 for thread in threads: 430 thread.start() 431 # Wait for the process to complete 432 if WaitForSingleObject(self.piProcInfo.hProcess, INFINITE) == \ 433 WAIT_FAILED: 434 raise ctypes.WinError() 435 # Wait for the I/O threads to complete 436 for thread in threads: 437 thread.join() 438 439 # Wait for the stdin thread to complete 440 if stdin_thread is not None: 441 stdin_thread.join() 442 443 def _stdin_raw_nonblock(self): 444 """Use the raw Win32 handle of sys.stdin to do non-blocking reads""" 445 # WARNING: This is experimental, and produces inconsistent results. 446 # It's possible for the handle not to be appropriate for use 447 # with WaitForSingleObject, among other things. 448 handle = msvcrt.get_osfhandle(sys.stdin.fileno()) 449 result = WaitForSingleObject(handle, 100) 450 if result == WAIT_FAILED: 451 raise ctypes.WinError() 452 elif result == WAIT_TIMEOUT: 453 print(".", end='') 454 return None 455 else: 456 data = ctypes.create_string_buffer(256) 457 bytesRead = DWORD(0) 458 print('?', end='') 459 460 if not ReadFile(handle, data, 256, 461 ctypes.byref(bytesRead), None): 462 raise ctypes.WinError() 463 # This ensures the non-blocking works with an actual console 464 # Not checking the error, so the processing will still work with 465 # other handle types 466 FlushConsoleInputBuffer(handle) 467 468 data = data.value 469 data = data.replace('\r\n', '\n') 470 data = data.replace('\r', '\n') 471 print(repr(data) + " ", end='') 472 return data 473 474 def _stdin_raw_block(self): 475 """Use a blocking stdin read""" 476 # The big problem with the blocking read is that it doesn't 477 # exit when it's supposed to in all contexts. An extra 478 # key-press may be required to trigger the exit. 479 try: 480 data = sys.stdin.read(1) 481 data = data.replace('\r', '\n') 482 return data 483 except WindowsError as we: 484 if we.winerror == ERROR_NO_DATA: 485 # This error occurs when the pipe is closed 486 return None 487 else: 488 # Otherwise let the error propagate 489 raise we 490 491 def _stdout_raw(self, s): 492 """Writes the string to stdout""" 493 print(s, end='', file=sys.stdout) 494 sys.stdout.flush() 495 496 def _stderr_raw(self, s): 497 """Writes the string to stdout""" 498 print(s, end='', file=sys.stderr) 499 sys.stderr.flush() 500 501 def _run_stdio(self): 502 """Runs the process using the system standard I/O. 503 504 IMPORTANT: stdin needs to be asynchronous, so the Python 505 sys.stdin object is not used. Instead, 506 msvcrt.kbhit/getwch are used asynchronously. 507 """ 508 # Disable Line and Echo mode 509 #lpMode = DWORD() 510 #handle = msvcrt.get_osfhandle(sys.stdin.fileno()) 511 #if GetConsoleMode(handle, ctypes.byref(lpMode)): 512 # set_console_mode = True 513 # if not SetConsoleMode(handle, lpMode.value & 514 # ~(ENABLE_ECHO_INPUT | ENABLE_LINE_INPUT | ENABLE_PROCESSED_INPUT)): 515 # raise ctypes.WinError() 516 517 if self.mergeout: 518 return self.run(stdout_func = self._stdout_raw, 519 stdin_func = self._stdin_raw_block) 520 else: 521 return self.run(stdout_func = self._stdout_raw, 522 stdin_func = self._stdin_raw_block, 523 stderr_func = self._stderr_raw) 524 525 # Restore the previous console mode 526 #if set_console_mode: 527 # if not SetConsoleMode(handle, lpMode.value): 528 # raise ctypes.WinError() 529 530 def __exit__(self, exc_type, exc_value, traceback): 531 if self.hstdin: 532 CloseHandle(self.hstdin) 533 self.hstdin = None 534 if self.hstdout: 535 CloseHandle(self.hstdout) 536 self.hstdout = None 537 if self.hstderr: 538 CloseHandle(self.hstderr) 539 self.hstderr = None 540 if self.piProcInfo is not None: 541 CloseHandle(self.piProcInfo.hProcess) 542 CloseHandle(self.piProcInfo.hThread) 543 self.piProcInfo = None 544 545 546 def system(cmd): 547 """Win32 version of os.system() that works with network shares. 548 549 Note that this implementation returns None, as meant for use in IPython. 550 551 Parameters 552 ---------- 553 cmd : str 554 A command to be executed in the system shell. 555 556 Returns 557 ------- 558 None : we explicitly do NOT return the subprocess status code, as this 559 utility is meant to be used extensively in IPython, where any return value 560 would trigger : func:`sys.displayhook` calls. 561 """ 562 with AvoidUNCPath() as path: 563 if path is not None: 564 cmd = '"pushd %s &&"%s' % (path, cmd) 565 with Win32ShellCommandController(cmd) as scc: 566 scc.run() 567 568 569 if __name__ == "__main__": 570 print("Test starting!") 571 #system("cmd") 572 system("python -i") 573 print("Test finished!") 574