1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5# A module to expose various thread/process/job related structures and
6# methods from kernel32
7#
8# The MIT License
9#
10# Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
11#
12# Additions and modifications written by Benjamin Smedberg
13# <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
14# <http://www.mozilla.org/>
15#
16# More Modifications
17# Copyright (c) 2006-2007 by Mike Taylor <bear@code-bear.com>
18# Copyright (c) 2007-2008 by Mikeal Rogers <mikeal@mozilla.com>
19#
20# By obtaining, using, and/or copying this software and/or its
21# associated documentation, you agree that you have read, understood,
22# and will comply with the following terms and conditions:
23#
24# Permission to use, copy, modify, and distribute this software and
25# its associated documentation for any purpose and without fee is
26# hereby granted, provided that the above copyright notice appears in
27# all copies, and that both that copyright notice and this permission
28# notice appear in supporting documentation, and that the name of the
29# author not be used in advertising or publicity pertaining to
30# distribution of the software without specific, written prior
31# permission.
32#
33# THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
34# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
35# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
36# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
37# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
38# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
39# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
40
41from __future__ import absolute_import, unicode_literals, print_function
42
43import subprocess
44import sys
45from ctypes import (
46    cast,
47    create_unicode_buffer,
48    c_ulong,
49    c_void_p,
50    POINTER,
51    sizeof,
52    Structure,
53    windll,
54    WinError,
55    WINFUNCTYPE,
56)
57from ctypes.wintypes import BOOL, BYTE, DWORD, HANDLE, LPCWSTR, LPWSTR, UINT, WORD
58
59from .qijo import QueryInformationJobObject
60
61LPVOID = c_void_p
62LPBYTE = POINTER(BYTE)
63LPDWORD = POINTER(DWORD)
64LPBOOL = POINTER(BOOL)
65LPULONG = POINTER(c_ulong)
66
67
68def ErrCheckBool(result, func, args):
69    """errcheck function for Windows functions that return a BOOL True
70    on success"""
71    if not result:
72        raise WinError()
73    return args
74
75
76# AutoHANDLE
77
78
79class AutoHANDLE(HANDLE):
80    """Subclass of HANDLE which will call CloseHandle() on deletion."""
81
82    CloseHandleProto = WINFUNCTYPE(BOOL, HANDLE)
83    CloseHandle = CloseHandleProto(("CloseHandle", windll.kernel32))
84    CloseHandle.errcheck = ErrCheckBool
85
86    def Close(self):
87        if self.value and self.value != HANDLE(-1).value:
88            self.CloseHandle(self)
89            self.value = 0
90
91    def __del__(self):
92        self.Close()
93
94    def __int__(self):
95        return self.value
96
97
98def ErrCheckHandle(result, func, args):
99    """errcheck function for Windows functions that return a HANDLE."""
100    if not result:
101        raise WinError()
102    return AutoHANDLE(result)
103
104
105# PROCESS_INFORMATION structure
106
107
108class PROCESS_INFORMATION(Structure):
109    _fields_ = [
110        ("hProcess", HANDLE),
111        ("hThread", HANDLE),
112        ("dwProcessID", DWORD),
113        ("dwThreadID", DWORD),
114    ]
115
116    def __init__(self):
117        Structure.__init__(self)
118
119        self.cb = sizeof(self)
120
121
122LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION)
123
124
125# STARTUPINFO structure
126
127
128class STARTUPINFO(Structure):
129    _fields_ = [
130        ("cb", DWORD),
131        ("lpReserved", LPWSTR),
132        ("lpDesktop", LPWSTR),
133        ("lpTitle", LPWSTR),
134        ("dwX", DWORD),
135        ("dwY", DWORD),
136        ("dwXSize", DWORD),
137        ("dwYSize", DWORD),
138        ("dwXCountChars", DWORD),
139        ("dwYCountChars", DWORD),
140        ("dwFillAttribute", DWORD),
141        ("dwFlags", DWORD),
142        ("wShowWindow", WORD),
143        ("cbReserved2", WORD),
144        ("lpReserved2", LPBYTE),
145        ("hStdInput", HANDLE),
146        ("hStdOutput", HANDLE),
147        ("hStdError", HANDLE),
148    ]
149
150
151LPSTARTUPINFO = POINTER(STARTUPINFO)
152
153SW_HIDE = 0
154
155STARTF_USESHOWWINDOW = 0x01
156STARTF_USESIZE = 0x02
157STARTF_USEPOSITION = 0x04
158STARTF_USECOUNTCHARS = 0x08
159STARTF_USEFILLATTRIBUTE = 0x10
160STARTF_RUNFULLSCREEN = 0x20
161STARTF_FORCEONFEEDBACK = 0x40
162STARTF_FORCEOFFFEEDBACK = 0x80
163STARTF_USESTDHANDLES = 0x100
164
165
166# EnvironmentBlock
167
168
169class EnvironmentBlock:
170    """An object which can be passed as the lpEnv parameter of CreateProcess.
171    It is initialized with a dictionary."""
172
173    def __init__(self, env):
174        if not env:
175            self._as_parameter_ = None
176        else:
177            values = []
178            fs_encoding = sys.getfilesystemencoding() or "mbcs"
179            for k, v in env.items():
180                if isinstance(k, bytes):
181                    k = k.decode(fs_encoding, "replace")
182                if isinstance(v, bytes):
183                    v = v.decode(fs_encoding, "replace")
184                values.append("{}={}".format(k, v))
185
186            # The lpEnvironment parameter of the 'CreateProcess' function expects a series
187            # of null terminated strings followed by a final null terminator. We write this
188            # value to a buffer and then cast it to LPCWSTR to avoid a Python ctypes bug
189            # that probihits embedded null characters (https://bugs.python.org/issue32745).
190            values = create_unicode_buffer("\0".join(values) + "\0")
191            self._as_parameter_ = cast(values, LPCWSTR)
192
193
194# Error Messages we need to watch for go here
195
196# https://msdn.microsoft.com/en-us/library/windows/desktop/ms681382(v=vs.85).aspx (0 - 499)
197ERROR_ACCESS_DENIED = 5
198ERROR_INVALID_PARAMETER = 87
199
200# http://msdn.microsoft.com/en-us/library/ms681388%28v=vs.85%29.aspx (500 - 999)
201ERROR_ABANDONED_WAIT_0 = 735
202
203# GetLastError()
204GetLastErrorProto = WINFUNCTYPE(DWORD)  # Return Type
205GetLastErrorFlags = ()
206GetLastError = GetLastErrorProto(("GetLastError", windll.kernel32), GetLastErrorFlags)
207
208# CreateProcess()
209
210CreateProcessProto = WINFUNCTYPE(
211    BOOL,  # Return type
212    LPCWSTR,  # lpApplicationName
213    LPWSTR,  # lpCommandLine
214    LPVOID,  # lpProcessAttributes
215    LPVOID,  # lpThreadAttributes
216    BOOL,  # bInheritHandles
217    DWORD,  # dwCreationFlags
218    LPVOID,  # lpEnvironment
219    LPCWSTR,  # lpCurrentDirectory
220    LPSTARTUPINFO,  # lpStartupInfo
221    LPPROCESS_INFORMATION,  # lpProcessInformation
222)
223
224CreateProcessFlags = (
225    (1, "lpApplicationName", None),
226    (1, "lpCommandLine"),
227    (1, "lpProcessAttributes", None),
228    (1, "lpThreadAttributes", None),
229    (1, "bInheritHandles", True),
230    (1, "dwCreationFlags", 0),
231    (1, "lpEnvironment", None),
232    (1, "lpCurrentDirectory", None),
233    (1, "lpStartupInfo"),
234    (2, "lpProcessInformation"),
235)
236
237
238def ErrCheckCreateProcess(result, func, args):
239    ErrCheckBool(result, func, args)
240    # return a tuple (hProcess, hThread, dwProcessID, dwThreadID)
241    pi = args[9]
242    return (
243        AutoHANDLE(pi.hProcess),
244        AutoHANDLE(pi.hThread),
245        pi.dwProcessID,
246        pi.dwThreadID,
247    )
248
249
250CreateProcess = CreateProcessProto(
251    ("CreateProcessW", windll.kernel32), CreateProcessFlags
252)
253CreateProcess.errcheck = ErrCheckCreateProcess
254
255# flags for CreateProcess
256CREATE_BREAKAWAY_FROM_JOB = 0x01000000
257CREATE_DEFAULT_ERROR_MODE = 0x04000000
258CREATE_NEW_CONSOLE = 0x00000010
259CREATE_NEW_PROCESS_GROUP = 0x00000200
260CREATE_NO_WINDOW = 0x08000000
261CREATE_SUSPENDED = 0x00000004
262CREATE_UNICODE_ENVIRONMENT = 0x00000400
263
264# Flags for IOCompletion ports (some of these would probably be defined if
265# we used the win32 extensions for python, but we don't want to do that if we
266# can help it.
267INVALID_HANDLE_VALUE = HANDLE(-1)  # From winbase.h
268
269# Self Defined Constants for IOPort <--> Job Object communication
270COMPKEY_TERMINATE = c_ulong(0)
271COMPKEY_JOBOBJECT = c_ulong(1)
272
273# flags for job limit information
274# see http://msdn.microsoft.com/en-us/library/ms684147%28VS.85%29.aspx
275JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x00000800
276JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x00001000
277JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000
278
279# Flags for Job Object Completion Port Message IDs from winnt.h
280# See also: http://msdn.microsoft.com/en-us/library/ms684141%28v=vs.85%29.aspx
281JOB_OBJECT_MSG_END_OF_JOB_TIME = 1
282JOB_OBJECT_MSG_END_OF_PROCESS_TIME = 2
283JOB_OBJECT_MSG_ACTIVE_PROCESS_LIMIT = 3
284JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO = 4
285JOB_OBJECT_MSG_NEW_PROCESS = 6
286JOB_OBJECT_MSG_EXIT_PROCESS = 7
287JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS = 8
288JOB_OBJECT_MSG_PROCESS_MEMORY_LIMIT = 9
289JOB_OBJECT_MSG_JOB_MEMORY_LIMIT = 10
290
291# See winbase.h
292DEBUG_ONLY_THIS_PROCESS = 0x00000002
293DEBUG_PROCESS = 0x00000001
294DETACHED_PROCESS = 0x00000008
295
296# OpenProcess -
297# https://msdn.microsoft.com/en-us/library/windows/desktop/ms684320(v=vs.85).aspx
298PROCESS_QUERY_INFORMATION = 0x0400
299PROCESS_VM_READ = 0x0010
300
301OpenProcessProto = WINFUNCTYPE(
302    HANDLE,  # Return type
303    DWORD,  # dwDesiredAccess
304    BOOL,  # bInheritHandle
305    DWORD,  # dwProcessId
306)
307
308OpenProcessFlags = (
309    (1, "dwDesiredAccess", 0),
310    (1, "bInheritHandle", False),
311    (1, "dwProcessId", 0),
312)
313
314
315def ErrCheckOpenProcess(result, func, args):
316    ErrCheckBool(result, func, args)
317
318    return AutoHANDLE(result)
319
320
321OpenProcess = OpenProcessProto(("OpenProcess", windll.kernel32), OpenProcessFlags)
322OpenProcess.errcheck = ErrCheckOpenProcess
323
324# GetQueuedCompletionPortStatus -
325# http://msdn.microsoft.com/en-us/library/aa364986%28v=vs.85%29.aspx
326GetQueuedCompletionStatusProto = WINFUNCTYPE(
327    BOOL,  # Return Type
328    HANDLE,  # Completion Port
329    LPDWORD,  # Msg ID
330    LPULONG,  # Completion Key
331    # PID Returned from the call (may be null)
332    LPULONG,
333    DWORD,
334)  # milliseconds to wait
335GetQueuedCompletionStatusFlags = (
336    (1, "CompletionPort", INVALID_HANDLE_VALUE),
337    (1, "lpNumberOfBytes", None),
338    (1, "lpCompletionKey", None),
339    (1, "lpPID", None),
340    (1, "dwMilliseconds", 0),
341)
342GetQueuedCompletionStatus = GetQueuedCompletionStatusProto(
343    ("GetQueuedCompletionStatus", windll.kernel32), GetQueuedCompletionStatusFlags
344)
345
346# CreateIOCompletionPort
347# Note that the completion key is just a number, not a pointer.
348CreateIoCompletionPortProto = WINFUNCTYPE(
349    HANDLE,  # Return Type
350    HANDLE,  # File Handle
351    HANDLE,  # Existing Completion Port
352    c_ulong,  # Completion Key
353    DWORD,
354)  # Number of Threads
355
356CreateIoCompletionPortFlags = (
357    (1, "FileHandle", INVALID_HANDLE_VALUE),
358    (1, "ExistingCompletionPort", 0),
359    (1, "CompletionKey", c_ulong(0)),
360    (1, "NumberOfConcurrentThreads", 0),
361)
362CreateIoCompletionPort = CreateIoCompletionPortProto(
363    ("CreateIoCompletionPort", windll.kernel32), CreateIoCompletionPortFlags
364)
365CreateIoCompletionPort.errcheck = ErrCheckHandle
366
367# SetInformationJobObject
368SetInformationJobObjectProto = WINFUNCTYPE(
369    BOOL,  # Return Type
370    HANDLE,  # Job Handle
371    DWORD,  # Type of Class next param is
372    LPVOID,  # Job Object Class
373    DWORD,
374)  # Job Object Class Length
375
376SetInformationJobObjectProtoFlags = (
377    (1, "hJob", None),
378    (1, "JobObjectInfoClass", None),
379    (1, "lpJobObjectInfo", None),
380    (1, "cbJobObjectInfoLength", 0),
381)
382SetInformationJobObject = SetInformationJobObjectProto(
383    ("SetInformationJobObject", windll.kernel32), SetInformationJobObjectProtoFlags
384)
385SetInformationJobObject.errcheck = ErrCheckBool
386
387# CreateJobObject()
388CreateJobObjectProto = WINFUNCTYPE(
389    HANDLE, LPVOID, LPCWSTR  # Return type  # lpJobAttributes  # lpName
390)
391
392CreateJobObjectFlags = ((1, "lpJobAttributes", None), (1, "lpName", None))
393
394CreateJobObject = CreateJobObjectProto(
395    ("CreateJobObjectW", windll.kernel32), CreateJobObjectFlags
396)
397CreateJobObject.errcheck = ErrCheckHandle
398
399# AssignProcessToJobObject()
400
401AssignProcessToJobObjectProto = WINFUNCTYPE(
402    BOOL, HANDLE, HANDLE  # Return type  # hJob  # hProcess
403)
404AssignProcessToJobObjectFlags = ((1, "hJob"), (1, "hProcess"))
405AssignProcessToJobObject = AssignProcessToJobObjectProto(
406    ("AssignProcessToJobObject", windll.kernel32), AssignProcessToJobObjectFlags
407)
408AssignProcessToJobObject.errcheck = ErrCheckBool
409
410# GetCurrentProcess()
411# because os.getPid() is way too easy
412GetCurrentProcessProto = WINFUNCTYPE(HANDLE)  # Return type
413GetCurrentProcessFlags = ()
414GetCurrentProcess = GetCurrentProcessProto(
415    ("GetCurrentProcess", windll.kernel32), GetCurrentProcessFlags
416)
417GetCurrentProcess.errcheck = ErrCheckHandle
418
419# IsProcessInJob()
420try:
421    IsProcessInJobProto = WINFUNCTYPE(
422        BOOL,  # Return type
423        HANDLE,  # Process Handle
424        HANDLE,  # Job Handle
425        LPBOOL,  # Result
426    )
427    IsProcessInJobFlags = (
428        (1, "ProcessHandle"),
429        (1, "JobHandle", HANDLE(0)),
430        (2, "Result"),
431    )
432    IsProcessInJob = IsProcessInJobProto(
433        ("IsProcessInJob", windll.kernel32), IsProcessInJobFlags
434    )
435    IsProcessInJob.errcheck = ErrCheckBool
436except AttributeError:
437    # windows 2k doesn't have this API
438    def IsProcessInJob(process):
439        return False
440
441
442# ResumeThread()
443
444
445def ErrCheckResumeThread(result, func, args):
446    if result == -1:
447        raise WinError()
448
449    return args
450
451
452ResumeThreadProto = WINFUNCTYPE(DWORD, HANDLE)  # Return type  # hThread
453ResumeThreadFlags = ((1, "hThread"),)
454ResumeThread = ResumeThreadProto(("ResumeThread", windll.kernel32), ResumeThreadFlags)
455ResumeThread.errcheck = ErrCheckResumeThread
456
457# TerminateProcess()
458
459TerminateProcessProto = WINFUNCTYPE(
460    BOOL, HANDLE, UINT  # Return type  # hProcess  # uExitCode
461)
462TerminateProcessFlags = ((1, "hProcess"), (1, "uExitCode", 127))
463TerminateProcess = TerminateProcessProto(
464    ("TerminateProcess", windll.kernel32), TerminateProcessFlags
465)
466TerminateProcess.errcheck = ErrCheckBool
467
468# TerminateJobObject()
469
470TerminateJobObjectProto = WINFUNCTYPE(
471    BOOL, HANDLE, UINT  # Return type  # hJob  # uExitCode
472)
473TerminateJobObjectFlags = ((1, "hJob"), (1, "uExitCode", 127))
474TerminateJobObject = TerminateJobObjectProto(
475    ("TerminateJobObject", windll.kernel32), TerminateJobObjectFlags
476)
477TerminateJobObject.errcheck = ErrCheckBool
478
479# WaitForSingleObject()
480
481WaitForSingleObjectProto = WINFUNCTYPE(
482    DWORD,
483    HANDLE,
484    DWORD,  # Return type  # hHandle  # dwMilliseconds
485)
486WaitForSingleObjectFlags = ((1, "hHandle"), (1, "dwMilliseconds", -1))
487WaitForSingleObject = WaitForSingleObjectProto(
488    ("WaitForSingleObject", windll.kernel32), WaitForSingleObjectFlags
489)
490
491# http://msdn.microsoft.com/en-us/library/ms681381%28v=vs.85%29.aspx
492INFINITE = -1
493WAIT_TIMEOUT = 0x0102
494WAIT_OBJECT_0 = 0x0
495WAIT_ABANDONED = 0x0080
496
497# http://msdn.microsoft.com/en-us/library/ms683189%28VS.85%29.aspx
498STILL_ACTIVE = 259
499
500# Used when we terminate a process.
501ERROR_CONTROL_C_EXIT = 0x23C
502ERROR_CONTROL_C_EXIT_DECIMAL = 3221225786
503
504# GetExitCodeProcess()
505
506GetExitCodeProcessProto = WINFUNCTYPE(
507    BOOL,
508    HANDLE,
509    LPDWORD,  # Return type  # hProcess  # lpExitCode
510)
511GetExitCodeProcessFlags = ((1, "hProcess"), (2, "lpExitCode"))
512GetExitCodeProcess = GetExitCodeProcessProto(
513    ("GetExitCodeProcess", windll.kernel32), GetExitCodeProcessFlags
514)
515GetExitCodeProcess.errcheck = ErrCheckBool
516
517
518def CanCreateJobObject():
519    currentProc = GetCurrentProcess()
520    if IsProcessInJob(currentProc):
521        jobinfo = QueryInformationJobObject(
522            HANDLE(0), "JobObjectExtendedLimitInformation"
523        )
524        limitflags = jobinfo["BasicLimitInformation"]["LimitFlags"]
525        return bool(limitflags & JOB_OBJECT_LIMIT_BREAKAWAY_OK) or bool(
526            limitflags & JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK
527        )
528    else:
529        return True
530
531
532# testing functions
533
534
535def parent():
536    print("Starting parent")
537    currentProc = GetCurrentProcess()
538    if IsProcessInJob(currentProc):
539        print("You should not be in a job object to test", file=sys.stderr)
540        sys.exit(1)
541    assert CanCreateJobObject()
542    print("File: %s" % __file__)
543    command = [sys.executable, __file__, "-child"]
544    print("Running command: %s" % command)
545    process = subprocess.Popen(command)
546    process.kill()
547    code = process.returncode
548    print("Child code: %s" % code)
549    assert code == 127
550
551
552def child():
553    print("Starting child")
554    currentProc = GetCurrentProcess()
555    injob = IsProcessInJob(currentProc)
556    print("Is in a job?: %s" % injob)
557    can_create = CanCreateJobObject()
558    print("Can create job?: %s" % can_create)
559    process = subprocess.Popen("c:\\windows\\notepad.exe")
560    assert process._job
561    jobinfo = QueryInformationJobObject(
562        process._job, "JobObjectExtendedLimitInformation"
563    )
564    print("Job info: %s" % jobinfo)
565    limitflags = jobinfo["BasicLimitInformation"]["LimitFlags"]
566    print("LimitFlags: %s" % limitflags)
567    process.kill()
568