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