michael@0: # A module to expose various thread/process/job related structures and
michael@0: # methods from kernel32
michael@0: #
michael@0: # The MIT License
michael@0: #
michael@0: # Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
michael@0: #
michael@0: # Additions and modifications written by Benjamin Smedberg
michael@0: # <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
michael@0: # <http://www.mozilla.org/>
michael@0: #
michael@0: # More Modifications
michael@0: # Copyright (c) 2006-2007 by Mike Taylor <bear@code-bear.com>
michael@0: # Copyright (c) 2007-2008 by Mikeal Rogers <mikeal@mozilla.com>
michael@0: #
michael@0: # By obtaining, using, and/or copying this software and/or its
michael@0: # associated documentation, you agree that you have read, understood,
michael@0: # and will comply with the following terms and conditions:
michael@0: #
michael@0: # Permission to use, copy, modify, and distribute this software and
michael@0: # its associated documentation for any purpose and without fee is
michael@0: # hereby granted, provided that the above copyright notice appears in
michael@0: # all copies, and that both that copyright notice and this permission
michael@0: # notice appear in supporting documentation, and that the name of the
michael@0: # author not be used in advertising or publicity pertaining to
michael@0: # distribution of the software without specific, written prior
michael@0: # permission.
michael@0: #
michael@0: # THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
michael@0: # INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
michael@0: # IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
michael@0: # CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
michael@0: # OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
michael@0: # NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
michael@0: # WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
michael@0: 
michael@0: from ctypes import c_void_p, POINTER, sizeof, Structure, windll, WinError, WINFUNCTYPE
michael@0: from ctypes.wintypes import BOOL, BYTE, DWORD, HANDLE, LPCWSTR, LPWSTR, UINT, WORD, \
michael@0:                             c_buffer, c_ulong, byref
michael@0: from qijo import QueryInformationJobObject
michael@0: 
michael@0: LPVOID = c_void_p
michael@0: LPBYTE = POINTER(BYTE)
michael@0: LPDWORD = POINTER(DWORD)
michael@0: LPBOOL = POINTER(BOOL)
michael@0: 
michael@0: def ErrCheckBool(result, func, args):
michael@0:     """errcheck function for Windows functions that return a BOOL True
michael@0:     on success"""
michael@0:     if not result:
michael@0:         raise WinError()
michael@0:     return args
michael@0: 
michael@0: 
michael@0: # AutoHANDLE
michael@0: 
michael@0: class AutoHANDLE(HANDLE):
michael@0:     """Subclass of HANDLE which will call CloseHandle() on deletion."""
michael@0:     
michael@0:     CloseHandleProto = WINFUNCTYPE(BOOL, HANDLE)
michael@0:     CloseHandle = CloseHandleProto(("CloseHandle", windll.kernel32))
michael@0:     CloseHandle.errcheck = ErrCheckBool
michael@0:     
michael@0:     def Close(self):
michael@0:         if self.value and self.value != HANDLE(-1).value:
michael@0:             self.CloseHandle(self)
michael@0:             self.value = 0
michael@0:     
michael@0:     def __del__(self):
michael@0:         self.Close()
michael@0: 
michael@0:     def __int__(self):
michael@0:         return self.value
michael@0: 
michael@0: def ErrCheckHandle(result, func, args):
michael@0:     """errcheck function for Windows functions that return a HANDLE."""
michael@0:     if not result:
michael@0:         raise WinError()
michael@0:     return AutoHANDLE(result)
michael@0: 
michael@0: # PROCESS_INFORMATION structure
michael@0: 
michael@0: class PROCESS_INFORMATION(Structure):
michael@0:     _fields_ = [("hProcess", HANDLE),
michael@0:                 ("hThread", HANDLE),
michael@0:                 ("dwProcessID", DWORD),
michael@0:                 ("dwThreadID", DWORD)]
michael@0: 
michael@0:     def __init__(self):
michael@0:         Structure.__init__(self)
michael@0:         
michael@0:         self.cb = sizeof(self)
michael@0: 
michael@0: LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION)
michael@0: 
michael@0: # STARTUPINFO structure
michael@0: 
michael@0: class STARTUPINFO(Structure):
michael@0:     _fields_ = [("cb", DWORD),
michael@0:                 ("lpReserved", LPWSTR),
michael@0:                 ("lpDesktop", LPWSTR),
michael@0:                 ("lpTitle", LPWSTR),
michael@0:                 ("dwX", DWORD),
michael@0:                 ("dwY", DWORD),
michael@0:                 ("dwXSize", DWORD),
michael@0:                 ("dwYSize", DWORD),
michael@0:                 ("dwXCountChars", DWORD),
michael@0:                 ("dwYCountChars", DWORD),
michael@0:                 ("dwFillAttribute", DWORD),
michael@0:                 ("dwFlags", DWORD),
michael@0:                 ("wShowWindow", WORD),
michael@0:                 ("cbReserved2", WORD),
michael@0:                 ("lpReserved2", LPBYTE),
michael@0:                 ("hStdInput", HANDLE),
michael@0:                 ("hStdOutput", HANDLE),
michael@0:                 ("hStdError", HANDLE)
michael@0:                 ]
michael@0: LPSTARTUPINFO = POINTER(STARTUPINFO)
michael@0: 
michael@0: SW_HIDE                 = 0
michael@0: 
michael@0: STARTF_USESHOWWINDOW    = 0x01
michael@0: STARTF_USESIZE          = 0x02
michael@0: STARTF_USEPOSITION      = 0x04
michael@0: STARTF_USECOUNTCHARS    = 0x08
michael@0: STARTF_USEFILLATTRIBUTE = 0x10
michael@0: STARTF_RUNFULLSCREEN    = 0x20
michael@0: STARTF_FORCEONFEEDBACK  = 0x40
michael@0: STARTF_FORCEOFFFEEDBACK = 0x80
michael@0: STARTF_USESTDHANDLES    = 0x100
michael@0: 
michael@0: # EnvironmentBlock
michael@0: 
michael@0: class EnvironmentBlock:
michael@0:     """An object which can be passed as the lpEnv parameter of CreateProcess.
michael@0:     It is initialized with a dictionary."""
michael@0: 
michael@0:     def __init__(self, dict):
michael@0:         if not dict:
michael@0:             self._as_parameter_ = None
michael@0:         else:
michael@0:             values = ["%s=%s" % (key, value)
michael@0:                       for (key, value) in dict.iteritems()]
michael@0:             values.append("")
michael@0:             self._as_parameter_ = LPCWSTR("\0".join(values))
michael@0:         
michael@0: # CreateProcess()
michael@0: 
michael@0: CreateProcessProto = WINFUNCTYPE(BOOL,                  # Return type
michael@0:                                  LPCWSTR,               # lpApplicationName
michael@0:                                  LPWSTR,                # lpCommandLine
michael@0:                                  LPVOID,                # lpProcessAttributes
michael@0:                                  LPVOID,                # lpThreadAttributes
michael@0:                                  BOOL,                  # bInheritHandles
michael@0:                                  DWORD,                 # dwCreationFlags
michael@0:                                  LPVOID,                # lpEnvironment
michael@0:                                  LPCWSTR,               # lpCurrentDirectory
michael@0:                                  LPSTARTUPINFO,         # lpStartupInfo
michael@0:                                  LPPROCESS_INFORMATION  # lpProcessInformation
michael@0:                                  )
michael@0: 
michael@0: CreateProcessFlags = ((1, "lpApplicationName", None),
michael@0:                       (1, "lpCommandLine"),
michael@0:                       (1, "lpProcessAttributes", None),
michael@0:                       (1, "lpThreadAttributes", None),
michael@0:                       (1, "bInheritHandles", True),
michael@0:                       (1, "dwCreationFlags", 0),
michael@0:                       (1, "lpEnvironment", None),
michael@0:                       (1, "lpCurrentDirectory", None),
michael@0:                       (1, "lpStartupInfo"),
michael@0:                       (2, "lpProcessInformation"))
michael@0: 
michael@0: def ErrCheckCreateProcess(result, func, args):
michael@0:     ErrCheckBool(result, func, args)
michael@0:     # return a tuple (hProcess, hThread, dwProcessID, dwThreadID)
michael@0:     pi = args[9]
michael@0:     return AutoHANDLE(pi.hProcess), AutoHANDLE(pi.hThread), pi.dwProcessID, pi.dwThreadID
michael@0: 
michael@0: CreateProcess = CreateProcessProto(("CreateProcessW", windll.kernel32),
michael@0:                                    CreateProcessFlags)
michael@0: CreateProcess.errcheck = ErrCheckCreateProcess
michael@0: 
michael@0: # flags for CreateProcess
michael@0: CREATE_BREAKAWAY_FROM_JOB = 0x01000000
michael@0: CREATE_DEFAULT_ERROR_MODE = 0x04000000
michael@0: CREATE_NEW_CONSOLE = 0x00000010
michael@0: CREATE_NEW_PROCESS_GROUP = 0x00000200
michael@0: CREATE_NO_WINDOW = 0x08000000
michael@0: CREATE_SUSPENDED = 0x00000004
michael@0: CREATE_UNICODE_ENVIRONMENT = 0x00000400
michael@0: 
michael@0: # flags for job limit information
michael@0: # see http://msdn.microsoft.com/en-us/library/ms684147%28VS.85%29.aspx
michael@0: JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x00000800
michael@0: JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x00001000
michael@0: 
michael@0: # XXX these flags should be documented
michael@0: DEBUG_ONLY_THIS_PROCESS = 0x00000002
michael@0: DEBUG_PROCESS = 0x00000001
michael@0: DETACHED_PROCESS = 0x00000008
michael@0: 
michael@0: # CreateJobObject()
michael@0: 
michael@0: CreateJobObjectProto = WINFUNCTYPE(HANDLE,             # Return type
michael@0:                                    LPVOID,             # lpJobAttributes
michael@0:                                    LPCWSTR             # lpName
michael@0:                                    )
michael@0: 
michael@0: CreateJobObjectFlags = ((1, "lpJobAttributes", None),
michael@0:                         (1, "lpName", None))
michael@0: 
michael@0: CreateJobObject = CreateJobObjectProto(("CreateJobObjectW", windll.kernel32),
michael@0:                                        CreateJobObjectFlags)
michael@0: CreateJobObject.errcheck = ErrCheckHandle
michael@0: 
michael@0: # AssignProcessToJobObject()
michael@0: 
michael@0: AssignProcessToJobObjectProto = WINFUNCTYPE(BOOL,      # Return type
michael@0:                                             HANDLE,    # hJob
michael@0:                                             HANDLE     # hProcess
michael@0:                                             )
michael@0: AssignProcessToJobObjectFlags = ((1, "hJob"),
michael@0:                                  (1, "hProcess"))
michael@0: AssignProcessToJobObject = AssignProcessToJobObjectProto(
michael@0:     ("AssignProcessToJobObject", windll.kernel32),
michael@0:     AssignProcessToJobObjectFlags)
michael@0: AssignProcessToJobObject.errcheck = ErrCheckBool
michael@0: 
michael@0: # GetCurrentProcess()
michael@0: # because os.getPid() is way too easy
michael@0: GetCurrentProcessProto = WINFUNCTYPE(HANDLE    # Return type
michael@0:                                      )
michael@0: GetCurrentProcessFlags = ()
michael@0: GetCurrentProcess = GetCurrentProcessProto(
michael@0:     ("GetCurrentProcess", windll.kernel32),
michael@0:     GetCurrentProcessFlags)
michael@0: GetCurrentProcess.errcheck = ErrCheckHandle
michael@0: 
michael@0: # IsProcessInJob()
michael@0: try:
michael@0:     IsProcessInJobProto = WINFUNCTYPE(BOOL,     # Return type
michael@0:                                       HANDLE,   # Process Handle
michael@0:                                       HANDLE,   # Job Handle
michael@0:                                       LPBOOL      # Result
michael@0:                                       )
michael@0:     IsProcessInJobFlags = ((1, "ProcessHandle"),
michael@0:                            (1, "JobHandle", HANDLE(0)),
michael@0:                            (2, "Result"))
michael@0:     IsProcessInJob = IsProcessInJobProto(
michael@0:         ("IsProcessInJob", windll.kernel32),
michael@0:         IsProcessInJobFlags)
michael@0:     IsProcessInJob.errcheck = ErrCheckBool 
michael@0: except AttributeError:
michael@0:     # windows 2k doesn't have this API
michael@0:     def IsProcessInJob(process):
michael@0:         return False
michael@0: 
michael@0: 
michael@0: # ResumeThread()
michael@0: 
michael@0: def ErrCheckResumeThread(result, func, args):
michael@0:     if result == -1:
michael@0:         raise WinError()
michael@0: 
michael@0:     return args
michael@0: 
michael@0: ResumeThreadProto = WINFUNCTYPE(DWORD,      # Return type
michael@0:                                 HANDLE      # hThread
michael@0:                                 )
michael@0: ResumeThreadFlags = ((1, "hThread"),)
michael@0: ResumeThread = ResumeThreadProto(("ResumeThread", windll.kernel32),
michael@0:                                  ResumeThreadFlags)
michael@0: ResumeThread.errcheck = ErrCheckResumeThread
michael@0: 
michael@0: # TerminateProcess()
michael@0: 
michael@0: TerminateProcessProto = WINFUNCTYPE(BOOL,   # Return type
michael@0:                                     HANDLE, # hProcess
michael@0:                                     UINT    # uExitCode
michael@0:                                     )
michael@0: TerminateProcessFlags = ((1, "hProcess"),
michael@0:                          (1, "uExitCode", 127))
michael@0: TerminateProcess = TerminateProcessProto(
michael@0:     ("TerminateProcess", windll.kernel32),
michael@0:     TerminateProcessFlags)
michael@0: TerminateProcess.errcheck = ErrCheckBool
michael@0: 
michael@0: # TerminateJobObject()
michael@0: 
michael@0: TerminateJobObjectProto = WINFUNCTYPE(BOOL,   # Return type
michael@0:                                       HANDLE, # hJob
michael@0:                                       UINT    # uExitCode
michael@0:                                       )
michael@0: TerminateJobObjectFlags = ((1, "hJob"),
michael@0:                            (1, "uExitCode", 127))
michael@0: TerminateJobObject = TerminateJobObjectProto(
michael@0:     ("TerminateJobObject", windll.kernel32),
michael@0:     TerminateJobObjectFlags)
michael@0: TerminateJobObject.errcheck = ErrCheckBool
michael@0: 
michael@0: # WaitForSingleObject()
michael@0: 
michael@0: WaitForSingleObjectProto = WINFUNCTYPE(DWORD,  # Return type
michael@0:                                        HANDLE, # hHandle
michael@0:                                        DWORD,  # dwMilliseconds
michael@0:                                        )
michael@0: WaitForSingleObjectFlags = ((1, "hHandle"),
michael@0:                             (1, "dwMilliseconds", -1))
michael@0: WaitForSingleObject = WaitForSingleObjectProto(
michael@0:     ("WaitForSingleObject", windll.kernel32),
michael@0:     WaitForSingleObjectFlags)
michael@0: 
michael@0: INFINITE = -1
michael@0: WAIT_TIMEOUT = 0x0102
michael@0: WAIT_OBJECT_0 = 0x0
michael@0: WAIT_ABANDONED = 0x0080
michael@0: WAIT_FAILED = 0xFFFFFFFF
michael@0: 
michael@0: # GetExitCodeProcess()
michael@0: 
michael@0: GetExitCodeProcessProto = WINFUNCTYPE(BOOL,    # Return type
michael@0:                                       HANDLE,  # hProcess
michael@0:                                       LPDWORD, # lpExitCode
michael@0:                                       )
michael@0: GetExitCodeProcessFlags = ((1, "hProcess"),
michael@0:                            (2, "lpExitCode"))
michael@0: GetExitCodeProcess = GetExitCodeProcessProto(
michael@0:     ("GetExitCodeProcess", windll.kernel32),
michael@0:     GetExitCodeProcessFlags)
michael@0: GetExitCodeProcess.errcheck = ErrCheckBool
michael@0: 
michael@0: def CanCreateJobObject():
michael@0:     # Running firefox in a job (from cfx) hangs on sites using flash plugin
michael@0:     # so job creation is turned off for now. (see Bug 768651).
michael@0:     return False
michael@0: 
michael@0: ### testing functions
michael@0: 
michael@0: def parent():
michael@0:     print 'Starting parent'
michael@0:     currentProc = GetCurrentProcess()
michael@0:     if IsProcessInJob(currentProc):
michael@0:         print >> sys.stderr, "You should not be in a job object to test"
michael@0:         sys.exit(1)
michael@0:     assert CanCreateJobObject()
michael@0:     print 'File: %s' % __file__
michael@0:     command = [sys.executable, __file__, '-child']
michael@0:     print 'Running command: %s' % command
michael@0:     process = Popen(command)
michael@0:     process.kill()
michael@0:     code = process.returncode
michael@0:     print 'Child code: %s' % code
michael@0:     assert code == 127
michael@0:         
michael@0: def child():
michael@0:     print 'Starting child'
michael@0:     currentProc = GetCurrentProcess()
michael@0:     injob = IsProcessInJob(currentProc)
michael@0:     print "Is in a job?: %s" % injob
michael@0:     can_create = CanCreateJobObject()
michael@0:     print 'Can create job?: %s' % can_create
michael@0:     process = Popen('c:\\windows\\notepad.exe')
michael@0:     assert process._job
michael@0:     jobinfo = QueryInformationJobObject(process._job, 'JobObjectExtendedLimitInformation')
michael@0:     print 'Job info: %s' % jobinfo
michael@0:     limitflags = jobinfo['BasicLimitInformation']['LimitFlags']
michael@0:     print 'LimitFlags: %s' % limitflags
michael@0:     process.kill()
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     import sys
michael@0:     from killableprocess import Popen
michael@0:     nargs = len(sys.argv[1:])
michael@0:     if nargs:
michael@0:         if nargs != 1 or sys.argv[1] != '-child':
michael@0:             raise AssertionError('Wrong flags; run like `python /path/to/winprocess.py`')
michael@0:         child()
michael@0:     else:
michael@0:         parent()