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 file,
3# You can obtain one at http://mozilla.org/MPL/2.0/.
4
5from __future__ import absolute_import
6
7from ctypes import (
8    c_void_p,
9    POINTER,
10    sizeof,
11    Structure,
12    windll,
13    WinError,
14    WINFUNCTYPE,
15    addressof,
16    c_size_t,
17    c_ulong,
18)
19from ctypes.wintypes import BOOL, BYTE, DWORD, HANDLE, LARGE_INTEGER
20
21import six
22
23LPVOID = c_void_p
24LPDWORD = POINTER(DWORD)
25SIZE_T = c_size_t
26ULONG_PTR = POINTER(c_ulong)
27
28# A ULONGLONG is a 64-bit unsigned integer.
29# Thus there are 8 bytes in a ULONGLONG.
30# XXX why not import c_ulonglong ?
31ULONGLONG = BYTE * 8
32
33
34class IO_COUNTERS(Structure):
35    # The IO_COUNTERS struct is 6 ULONGLONGs.
36    # TODO: Replace with non-dummy fields.
37    _fields_ = [("dummy", ULONGLONG * 6)]
38
39
40class JOBOBJECT_BASIC_ACCOUNTING_INFORMATION(Structure):
41    _fields_ = [
42        ("TotalUserTime", LARGE_INTEGER),
43        ("TotalKernelTime", LARGE_INTEGER),
44        ("ThisPeriodTotalUserTime", LARGE_INTEGER),
45        ("ThisPeriodTotalKernelTime", LARGE_INTEGER),
46        ("TotalPageFaultCount", DWORD),
47        ("TotalProcesses", DWORD),
48        ("ActiveProcesses", DWORD),
49        ("TotalTerminatedProcesses", DWORD),
50    ]
51
52
53class JOBOBJECT_BASIC_AND_IO_ACCOUNTING_INFORMATION(Structure):
54    _fields_ = [
55        ("BasicInfo", JOBOBJECT_BASIC_ACCOUNTING_INFORMATION),
56        ("IoInfo", IO_COUNTERS),
57    ]
58
59
60# see http://msdn.microsoft.com/en-us/library/ms684147%28VS.85%29.aspx
61class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure):
62    _fields_ = [
63        ("PerProcessUserTimeLimit", LARGE_INTEGER),
64        ("PerJobUserTimeLimit", LARGE_INTEGER),
65        ("LimitFlags", DWORD),
66        ("MinimumWorkingSetSize", SIZE_T),
67        ("MaximumWorkingSetSize", SIZE_T),
68        ("ActiveProcessLimit", DWORD),
69        ("Affinity", ULONG_PTR),
70        ("PriorityClass", DWORD),
71        ("SchedulingClass", DWORD),
72    ]
73
74
75class JOBOBJECT_ASSOCIATE_COMPLETION_PORT(Structure):
76    _fields_ = [("CompletionKey", c_ulong), ("CompletionPort", HANDLE)]
77
78
79# see http://msdn.microsoft.com/en-us/library/ms684156%28VS.85%29.aspx
80class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure):
81    _fields_ = [
82        ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION),
83        ("IoInfo", IO_COUNTERS),
84        ("ProcessMemoryLimit", SIZE_T),
85        ("JobMemoryLimit", SIZE_T),
86        ("PeakProcessMemoryUsed", SIZE_T),
87        ("PeakJobMemoryUsed", SIZE_T),
88    ]
89
90
91# These numbers below come from:
92# http://msdn.microsoft.com/en-us/library/ms686216%28v=vs.85%29.aspx
93JobObjectAssociateCompletionPortInformation = 7
94JobObjectBasicAndIoAccountingInformation = 8
95JobObjectExtendedLimitInformation = 9
96
97
98class JobObjectInfo(object):
99    mapping = {
100        "JobObjectBasicAndIoAccountingInformation": 8,
101        "JobObjectExtendedLimitInformation": 9,
102        "JobObjectAssociateCompletionPortInformation": 7,
103    }
104    structures = {
105        7: JOBOBJECT_ASSOCIATE_COMPLETION_PORT,
106        8: JOBOBJECT_BASIC_AND_IO_ACCOUNTING_INFORMATION,
107        9: JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
108    }
109
110    def __init__(self, _class):
111        if isinstance(_class, six.string_types):
112            assert _class in self.mapping, "Class should be one of %s; you gave %s" % (
113                self.mapping,
114                _class,
115            )
116            _class = self.mapping[_class]
117        assert _class in self.structures, "Class should be one of %s; you gave %s" % (
118            self.structures,
119            _class,
120        )
121        self.code = _class
122        self.info = self.structures[_class]()
123
124
125QueryInformationJobObjectProto = WINFUNCTYPE(
126    BOOL,  # Return type
127    HANDLE,  # hJob
128    DWORD,  # JobObjectInfoClass
129    LPVOID,  # lpJobObjectInfo
130    DWORD,  # cbJobObjectInfoLength
131    LPDWORD,  # lpReturnLength
132)
133
134QueryInformationJobObjectFlags = (
135    (1, "hJob"),
136    (1, "JobObjectInfoClass"),
137    (1, "lpJobObjectInfo"),
138    (1, "cbJobObjectInfoLength"),
139    (1, "lpReturnLength", None),
140)
141
142_QueryInformationJobObject = QueryInformationJobObjectProto(
143    ("QueryInformationJobObject", windll.kernel32), QueryInformationJobObjectFlags
144)
145
146
147class SubscriptableReadOnlyStruct(object):
148    def __init__(self, struct):
149        self._struct = struct
150
151    def _delegate(self, name):
152        result = getattr(self._struct, name)
153        if isinstance(result, Structure):
154            return SubscriptableReadOnlyStruct(result)
155        return result
156
157    def __getitem__(self, name):
158        match = [fname for fname, ftype in self._struct._fields_ if fname == name]
159        if match:
160            return self._delegate(name)
161        raise KeyError(name)
162
163    def __getattr__(self, name):
164        return self._delegate(name)
165
166
167def QueryInformationJobObject(hJob, JobObjectInfoClass):
168    jobinfo = JobObjectInfo(JobObjectInfoClass)
169    result = _QueryInformationJobObject(
170        hJob=hJob,
171        JobObjectInfoClass=jobinfo.code,
172        lpJobObjectInfo=addressof(jobinfo.info),
173        cbJobObjectInfoLength=sizeof(jobinfo.info),
174    )
175    if not result:
176        raise WinError()
177    return SubscriptableReadOnlyStruct(jobinfo.info)
178