From 6a6c4e94ca1125e3bf07cc3437e7316097fc0f9d 2012-02-14 23:28:02 From: Mark Wiebe Date: 2012-02-14 23:28:02 Subject: [PATCH] Create a shell application controller using direct Win32 API --- diff --git a/IPython/utils/_process_win32_controller.py b/IPython/utils/_process_win32_controller.py new file mode 100644 index 0000000..93a163d --- /dev/null +++ b/IPython/utils/_process_win32_controller.py @@ -0,0 +1,502 @@ +"""Windows-specific implementation of process utilities with direct WinAPI. + +This file is meant to be used by process.py +""" + +#----------------------------------------------------------------------------- +# Copyright (C) 2010-2011 The IPython Development Team +# +# Distributed under the terms of the BSD License. The full license is in +# the file COPYING, distributed as part of this software. +#----------------------------------------------------------------------------- + +from __future__ import print_function + +# stdlib +import os, sys, time, threading +import ctypes, msvcrt + +# Win32 API types needed for the API calls +from ctypes import POINTER +from ctypes.wintypes import HANDLE, HLOCAL, LPVOID, WORD, DWORD, BOOL, \ + ULONG, LPCWSTR +LPDWORD = POINTER(DWORD) +LPHANDLE = POINTER(HANDLE) +ULONG_PTR = POINTER(ULONG) +class SECURITY_ATTRIBUTES(ctypes.Structure): + _fields_ = [("nLength", DWORD), + ("lpSecurityDescriptor", LPVOID), + ("bInheritHandle", BOOL)] +LPSECURITY_ATTRIBUTES = POINTER(SECURITY_ATTRIBUTES) +class STARTUPINFO(ctypes.Structure): + _fields_ = [("cb", DWORD), + ("lpReserved", LPCWSTR), + ("lpDesktop", LPCWSTR), + ("lpTitle", LPCWSTR), + ("dwX", DWORD), + ("dwY", DWORD), + ("dwXSize", DWORD), + ("dwYSize", DWORD), + ("dwXCountChars", DWORD), + ("dwYCountChars", DWORD), + ("dwFillAttribute", DWORD), + ("dwFlags", DWORD), + ("wShowWindow", WORD), + ("cbReserved2", WORD), + ("lpReserved2", LPVOID), + ("hStdInput", HANDLE), + ("hStdOutput", HANDLE), + ("hStdError", HANDLE)] +LPSTARTUPINFO = POINTER(STARTUPINFO) +class PROCESS_INFORMATION(ctypes.Structure): + _fields_ = [("hProcess", HANDLE), + ("hThread", HANDLE), + ("dwProcessId", DWORD), + ("dwThreadId", DWORD)] +LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION) + +# Win32 API constants needed +ERROR_HANDLE_EOF = 38 +ERROR_BROKEN_PIPE = 109 +HANDLE_FLAG_INHERIT = 0x0001 +STARTF_USESTDHANDLES = 0x0100 +CREATE_SUSPENDED = 0x0004 +CREATE_NEW_CONSOLE = 0x0010 +STILL_ACTIVE = 259 +WAIT_TIMEOUT = 0x0102 +WAIT_FAILED = 0xFFFFFFFF +INFINITE = 0xFFFFFFFF +DUPLICATE_SAME_ACCESS = 0x00000002 + +# Win32 API functions needed +GetLastError = ctypes.windll.kernel32.GetLastError +GetLastError.argtypes = [] +GetLastError.restype = DWORD + +CreateFile = ctypes.windll.kernel32.CreateFileW +CreateFile.argtypes = [LPCWSTR, DWORD, DWORD, LPVOID, DWORD, DWORD, HANDLE] +CreateFile.restype = HANDLE + +CreatePipe = ctypes.windll.kernel32.CreatePipe +CreatePipe.argtypes = [POINTER(HANDLE), POINTER(HANDLE), + LPSECURITY_ATTRIBUTES, DWORD] +CreatePipe.restype = BOOL + +CreateProcess = ctypes.windll.kernel32.CreateProcessW +CreateProcess.argtypes = [LPCWSTR, LPCWSTR, LPSECURITY_ATTRIBUTES, + LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFO, + LPPROCESS_INFORMATION] +CreateProcess.restype = BOOL + +GetExitCodeProcess = ctypes.windll.kernel32.GetExitCodeProcess +GetExitCodeProcess.argtypes = [HANDLE, LPDWORD] +GetExitCodeProcess.restype = BOOL + +GetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess +GetCurrentProcess.argtypes = [] +GetCurrentProcess.restype = HANDLE + +ResumeThread = ctypes.windll.kernel32.ResumeThread +ResumeThread.argtypes = [HANDLE] +ResumeThread.restype = DWORD + +ReadFile = ctypes.windll.kernel32.ReadFile +ReadFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID] +ReadFile.restype = BOOL + +WriteFile = ctypes.windll.kernel32.WriteFile +WriteFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID] +WriteFile.restype = BOOL + +WaitForSingleObject = ctypes.windll.kernel32.WaitForSingleObject +WaitForSingleObject.argtypes = [HANDLE, DWORD] +WaitForSingleObject.restype = DWORD + +DuplicateHandle = ctypes.windll.kernel32.DuplicateHandle +DuplicateHandle.argtypes = [HANDLE, HANDLE, HANDLE, LPHANDLE, + DWORD, BOOL, DWORD] +DuplicateHandle.restype = BOOL + +SetHandleInformation = ctypes.windll.kernel32.SetHandleInformation +SetHandleInformation.argtypes = [HANDLE, DWORD, DWORD] +SetHandleInformation.restype = BOOL + +CloseHandle = ctypes.windll.kernel32.CloseHandle +CloseHandle.argtypes = [HANDLE] +CloseHandle.restype = BOOL + +CommandLineToArgvW = ctypes.windll.shell32.CommandLineToArgvW +CommandLineToArgvW.argtypes = [LPCWSTR, POINTER(ctypes.c_int)] +CommandLineToArgvW.restype = POINTER(LPCWSTR) + +LocalFree = ctypes.windll.kernel32.LocalFree +LocalFree.argtypes = [HLOCAL] +LocalFree.restype = HLOCAL + +class AvoidUNCPath(object): + """A context manager to protect command execution from UNC paths. + + In the Win32 API, commands can't be invoked with the cwd being a UNC path. + This context manager temporarily changes directory to the 'C:' drive on + entering, and restores the original working directory on exit. + + The context manager returns the starting working directory *if* it made a + change and None otherwise, so that users can apply the necessary adjustment + to their system calls in the event of a change. + + Example + ------- + :: + cmd = 'dir' + with AvoidUNCPath() as path: + if path is not None: + cmd = '"pushd %s &&"%s' % (path, cmd) + os.system(cmd) + """ + def __enter__(self): + self.path = os.getcwdu() + self.is_unc_path = self.path.startswith(r"\\") + if self.is_unc_path: + # change to c drive (as cmd.exe cannot handle UNC addresses) + os.chdir("C:") + return self.path + else: + # We return None to signal that there was no change in the working + # directory + return None + + def __exit__(self, exc_type, exc_value, traceback): + if self.is_unc_path: + os.chdir(self.path) + + +class Win32ShellCommandController(object): + """Runs a shell command in a 'with' context. + + This implementation is Win32-specific. + + Example: + # Runs the command interactively with default console stdin/stdout + with ShellCommandController('python -i') as scc: + scc.run() + + # Runs the command using the provided functions for stdin/stdout + def my_stdout_func(s): + # print or save the string 's' + write_to_stdout(s) + def my_stdin_func(): + # If input is available, return it as a string. + if input_available(): + return get_input() + # If no input available, return None after a short delay to + # keep from blocking. + else: + time.sleep(0.01) + return None + + with ShellCommandController('python -i') as scc: + scc.run(my_stdout_func, my_stdin_func) + """ + + def __init__(self, cmd, mergeout = True): + """Initializes the shell command controller. + + The cmd is the program to execute, and mergeout is + whether to blend stdout and stderr into one output + in stdout. Merging them together in this fashion more + reliably keeps stdout and stderr in the correct order + especially for interactive shell usage. + """ + self.cmd = cmd + self.mergeout = mergeout + + def __enter__(self): + cmd = self.cmd + mergeout = self.mergeout + + self.hstdout, self.hstdin, self.hstderr = None, None, None + self.piProcInfo = None + try: + p_hstdout, c_hstdout, p_hstderr, \ + c_hstderr, p_hstdin, c_hstdin = [None]*6 + + # SECURITY_ATTRIBUTES with inherit handle set to True + saAttr = SECURITY_ATTRIBUTES() + saAttr.nLength = ctypes.sizeof(saAttr) + saAttr.bInheritHandle = True + saAttr.lpSecurityDescriptor = None + + def create_pipe(uninherit): + """Creates a Windows pipe, which consists of two handles. + + The 'uninherit' parameter controls which handle is not + inherited by the child process. + """ + handles = HANDLE(), HANDLE() + if not CreatePipe(ctypes.byref(handles[0]), + ctypes.byref(handles[1]), ctypes.byref(saAttr), 0): + raise ctypes.WinError() + if not SetHandleInformation(handles[uninherit], + HANDLE_FLAG_INHERIT, 0): + raise ctypes.WinError() + return handles[0].value, handles[1].value + + p_hstdout, c_hstdout = create_pipe(uninherit=0) + # 'mergeout' signals that stdout and stderr should be merged. + # We do that by using one pipe for both of them. + if mergeout: + c_hstderr = HANDLE() + if not DuplicateHandle(GetCurrentProcess(), c_hstdout, + GetCurrentProcess(), ctypes.byref(c_hstderr), + 0, True, DUPLICATE_SAME_ACCESS): + raise ctypes.WinError() + else: + p_hstderr, c_hstderr = create_pipe(uninherit=0) + c_hstdin, p_hstdin = create_pipe(uninherit=1) + + # Create the process object + piProcInfo = PROCESS_INFORMATION() + siStartInfo = STARTUPINFO() + siStartInfo.cb = ctypes.sizeof(siStartInfo) + siStartInfo.hStdInput = c_hstdin + siStartInfo.hStdOutput = c_hstdout + siStartInfo.hStdError = c_hstderr + siStartInfo.dwFlags = STARTF_USESTDHANDLES + dwCreationFlags = CREATE_SUSPENDED # | CREATE_NEW_CONSOLE + + if not CreateProcess(None, + u"cmd.exe /c " + cmd, + None, None, True, dwCreationFlags, + None, None, ctypes.byref(siStartInfo), + ctypes.byref(piProcInfo)): + raise ctypes.WinError() + + # Close this process's versions of the child handles + CloseHandle(c_hstdin) + c_hstdin = None + CloseHandle(c_hstdout) + c_hstdout = None + if c_hstderr != None: + CloseHandle(c_hstderr) + c_hstderr = None + + # Transfer ownership of the parent handles to the object + self.hstdin = p_hstdin + p_hstdin = None + self.hstdout = p_hstdout + p_hstdout = None + if not mergeout: + self.hstderr = p_hstderr + p_hstderr = None + self.piProcInfo = piProcInfo + + finally: + if p_hstdin: + CloseHandle(p_hstdin) + if c_hstdin: + CloseHandle(c_hstdin) + if p_hstdout: + CloseHandle(p_hstdout) + if c_hstdout: + CloseHandle(c_hstdout) + if p_hstderr: + CloseHandle(p_hstderr) + if c_hstderr: + CloseHandle(c_hstderr) + + return self + + def _stdin_thread(self, handle, hprocess, func, stdout_func): + # TODO: Use WaitForInputIdle to avoid calling func() until + # an input is actually requested. + exitCode = DWORD() + bytesWritten = DWORD(0) + while True: + #print("stdin thread loop start") + # Get the input string (may be bytes or unicode) + data = func() + + # None signals to poll whether the process has exited + if data is None: + #print("checking for process completion") + if not GetExitCodeProcess(hprocess, ctypes.byref(exitCode)): + raise ctypes.WinError() + if exitCode.value != STILL_ACTIVE: + return + # TESTING: Does zero-sized writefile help? + if not WriteFile(handle, "", 0, + ctypes.byref(bytesWritten), None): + raise ctypes.WinError() + continue + #print("\nGot str %s\n" % repr(data), file=sys.stderr) + + # Encode the string to the console encoding + if isinstance(data, unicode): #FIXME: Python3 + data = data.encode('utf_8') + + # What we have now must be a string of bytes + if not isinstance(data, str): #FIXME: Python3 + raise RuntimeError("internal stdin function string error") + + # An empty string signals EOF + if len(data) == 0: + return + + # In a windows console, sometimes the input is echoed, + # but sometimes not. How do we determine when to do this? + stdout_func(data) + # WriteFile may not accept all the data at once. + # Loop until everything is processed + while len(data) != 0: + #print("Calling writefile") + if not WriteFile(handle, data, len(data), + ctypes.byref(bytesWritten), None): + raise ctypes.WinError() + #print("Called writefile") + data = data[bytesWritten.value:] + + def _stdout_thread(self, handle, func): + # Allocate the output buffer + data = ctypes.create_string_buffer(4096) + while True: + bytesRead = DWORD(0) + if not ReadFile(handle, data, 4096, + ctypes.byref(bytesRead), None): + le = GetLastError() + if le == ERROR_BROKEN_PIPE: + return + else: + raise ctypes.WinError() + # FIXME: Python3 + s = data.value[0:bytesRead.value] + #print("\nv: %s" % repr(s), file=sys.stderr) + func(s.decode('utf_8', 'replace')) + + def run(self, stdout_func = None, stdin_func = None, stderr_func = None): + """Runs the process, using the provided functions for I/O. + + The function stdin_func should return strings whenever a + character or characters become available. + The functions stdout_func and stderr_func are called whenever + something is printed to stdout or stderr, respectively. + These functions are called from different threads (but not + concurrently, because of the GIL). + """ + if stdout_func == None and stdin_func == None and stderr_func == None: + return self._run_stdio() + + if stderr_func != None and self.hstderr == None: + raise RuntimeError("Shell command was initiated with " + "merged stdin/stdout, but a separate stderr_func " + "was provided to the run() method") + + # Create a thread for each input/output handle + threads = [] + if stdin_func: + threads.append(threading.Thread(target=self._stdin_thread, + args=(self.hstdin, self.piProcInfo.hProcess, + stdin_func, stdout_func))) + threads.append(threading.Thread(target=self._stdout_thread, + args=(self.hstdout, stdout_func))) + if self.hstderr != None: + if stderr_func == None: + stderr_func = stdout_func + threads.append(threading.Thread(target=self._stdout_thread, + args=(self.hstderr, stderr_func))) + # Start the I/O threads and the process + if ResumeThread(self.piProcInfo.hThread) == 0xFFFFFFFF: + raise ctypes.WinError() + for thread in threads: + thread.start() + # Wait for the process to complete + if WaitForSingleObject(self.piProcInfo.hProcess, INFINITE) == \ + WAIT_FAILED: + raise ctypes.WinError() + # Wait for the I/O threads to complete + for thread in threads: + thread.join() + + def _stdin_raw(self): + """Uses msvcrt.kbhit/getwch to do read stdin without blocking""" + if msvcrt.kbhit(): + #s = msvcrt.getwch() + s = msvcrt.getwch() + # Key code for Enter is '\r', but need to give back '\n' + if s == u'\r': + s = u'\n' + return s + else: + # This should make it poll at about 100 Hz, which + # is hopefully good enough to be responsive but + # doesn't waste CPU. + time.sleep(0.01) + return None + + def _stdout_raw(self, s): + """Writes the string to stdout""" + print(s, end='', file=sys.stdout) + + def _stderr_raw(self, s): + """Writes the string to stdout""" + print(s, end='', file=sys.stderr) + + def _run_stdio(self): + """Runs the process using the system standard I/O. + + IMPORTANT: stdin needs to be asynchronous, so the Python + sys.stdin object is not used. Instead, + msvcrt.kbhit/getwch are used asynchronously. + """ + if self.hstderr != None: + return self.run(stdout_func = self._stdout_raw, + stdin_func = self._stdin_raw, + stderr_func = self._stderr_raw) + else: + return self.run(stdout_func = self._stdout_raw, + stdin_func = self._stdin_raw) + + + def __exit__(self, exc_type, exc_value, traceback): + if self.hstdin: + CloseHandle(self.hstdin) + self.hstdin = None + if self.hstdout: + CloseHandle(self.hstdout) + self.hstdout = None + if self.hstderr: + CloseHandle(self.hstderr) + self.hstderr = None + if self.piProcInfo != None: + CloseHandle(self.piProcInfo.hProcess) + CloseHandle(self.piProcInfo.hThread) + self.piProcInfo = None + + +def system(cmd): + """Win32 version of os.system() that works with network shares. + + Note that this implementation returns None, as meant for use in IPython. + + Parameters + ---------- + cmd : str + A command to be executed in the system shell. + + Returns + ------- + None : we explicitly do NOT return the subprocess status code, as this + utility is meant to be used extensively in IPython, where any return value + would trigger :func:`sys.displayhook` calls. + """ + with AvoidUNCPath() as path: + if path is not None: + cmd = '"pushd %s &&"%s' % (path, cmd) + with Win32ShellCommandController(cmd) as scc: + scc.run() + + +if __name__ == "__main__": + print("Test starting!") + #system("cmd") + system("python -i") + print("Test finished!")