1 /*
2  * Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
3  * Use of this source code is governed by a BSD-style license that can be
4  * found in the LICENSE file.
5  *
6  * Windows platform-specific module methods for _psutil_windows.
7  *
8  * List of undocumented Windows NT APIs which are used in here and in
9  * other modules:
10  * - NtQuerySystemInformation
11  * - NtQueryInformationProcess
12  * - NtQueryObject
13  * - NtSuspendProcess
14  * - NtResumeProcess
15  */
16 
17 // Fixes clash between winsock2.h and windows.h
18 #define WIN32_LEAN_AND_MEAN
19 
20 #include <Python.h>
21 #include <windows.h>
22 #include <Psapi.h>
23 #include <signal.h>
24 #include <WinIoCtl.h>  // disk_io_counters()
25 #include <tchar.h>
26 #include <tlhelp32.h>
27 #include <wtsapi32.h>  // users()
28 #include <PowrProf.h>  // cpu_freq()
29 #if (_WIN32_WINNT >= 0x0600) // Windows >= Vista
30 #include <ws2tcpip.h>  // net_connections()
31 #endif
32 
33 // Link with Iphlpapi.lib
34 #pragma comment(lib, "IPHLPAPI.lib")
35 
36 #include "arch/windows/ntextapi.h"
37 #include "arch/windows/global.h"
38 #include "arch/windows/security.h"
39 #include "arch/windows/process_info.h"
40 #include "arch/windows/process_handles.h"
41 #include "arch/windows/inet_ntop.h"
42 #include "arch/windows/services.h"
43 #include "arch/windows/wmi.h"
44 #include "_psutil_common.h"
45 
46 
47 /*
48  * ============================================================================
49  * Utilities
50  * ============================================================================
51  */
52 
53 #define MALLOC(x) HeapAlloc(GetProcessHeap(), 0, (x))
54 #define FREE(x) HeapFree(GetProcessHeap(), 0, (x))
55 #define LO_T 1e-7
56 #define HI_T 429.4967296
57 #define BYTESWAP_USHORT(x) ((((USHORT)(x) << 8) | ((USHORT)(x) >> 8)) & 0xffff)
58 #ifndef AF_INET6
59 #define AF_INET6 23
60 #endif
61 
62 
63 PIP_ADAPTER_ADDRESSES
psutil_get_nic_addresses()64 psutil_get_nic_addresses() {
65     // allocate a 15 KB buffer to start with
66     int outBufLen = 15000;
67     DWORD dwRetVal = 0;
68     ULONG attempts = 0;
69     PIP_ADAPTER_ADDRESSES pAddresses = NULL;
70 
71     do {
72         pAddresses = (IP_ADAPTER_ADDRESSES *) malloc(outBufLen);
73         if (pAddresses == NULL) {
74             PyErr_NoMemory();
75             return NULL;
76         }
77 
78         dwRetVal = GetAdaptersAddresses(AF_UNSPEC, 0, NULL, pAddresses,
79                                         &outBufLen);
80         if (dwRetVal == ERROR_BUFFER_OVERFLOW) {
81             free(pAddresses);
82             pAddresses = NULL;
83         }
84         else {
85             break;
86         }
87 
88         attempts++;
89     } while ((dwRetVal == ERROR_BUFFER_OVERFLOW) && (attempts < 3));
90 
91     if (dwRetVal != NO_ERROR) {
92         PyErr_SetString(
93             PyExc_RuntimeError, "GetAdaptersAddresses() syscall failed.");
94         return NULL;
95     }
96 
97     return pAddresses;
98 }
99 
100 
101 /*
102  * Return the number of logical, active CPUs. Return 0 if undetermined.
103  * See discussion at: https://bugs.python.org/issue33166#msg314631
104  */
105 unsigned int
psutil_get_num_cpus(int fail_on_err)106 psutil_get_num_cpus(int fail_on_err) {
107     unsigned int ncpus = 0;
108 
109     // Minimum requirement: Windows 7
110     if (psutil_GetActiveProcessorCount != NULL) {
111         ncpus = psutil_GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
112         if ((ncpus == 0) && (fail_on_err == 1)) {
113             PyErr_SetFromWindowsErr(0);
114         }
115     }
116     else {
117         psutil_debug("GetActiveProcessorCount() not available; "
118                      "using GetSystemInfo()");
119         ncpus = (unsigned int)PSUTIL_SYSTEM_INFO.dwNumberOfProcessors;
120         if ((ncpus <= 0) && (fail_on_err == 1)) {
121             PyErr_SetString(
122                 PyExc_RuntimeError,
123                 "GetSystemInfo() failed to retrieve CPU count");
124         }
125     }
126     return ncpus;
127 }
128 
129 
130 /*
131  * ============================================================================
132  * Public Python API
133  * ============================================================================
134  */
135 
136 // Raised by Process.wait().
137 static PyObject *TimeoutExpired;
138 static PyObject *TimeoutAbandoned;
139 
140 /*
141  * Return a Python float representing the system uptime expressed in seconds
142  * since the epoch.
143  */
144 static PyObject *
psutil_boot_time(PyObject * self,PyObject * args)145 psutil_boot_time(PyObject *self, PyObject *args) {
146     ULONGLONG uptime;
147     time_t pt;
148     FILETIME fileTime;
149     ULONGLONG ll;
150 
151     GetSystemTimeAsFileTime(&fileTime);
152     /*
153     HUGE thanks to:
154     http://johnstewien.spaces.live.com/blog/cns!E6885DB5CEBABBC8!831.entry
155 
156     This function converts the FILETIME structure to the 32 bit
157     Unix time structure.
158     The time_t is a 32-bit value for the number of seconds since
159     January 1, 1970. A FILETIME is a 64-bit for the number of
160     100-nanosecond periods since January 1, 1601. Convert by
161     subtracting the number of 100-nanosecond period between 01-01-1970
162     and 01-01-1601, from time_t the divide by 1e+7 to get to the same
163     base granularity.
164     */
165     ll = (((ULONGLONG)
166         (fileTime.dwHighDateTime)) << 32) + fileTime.dwLowDateTime;
167     pt = (time_t)((ll - 116444736000000000ull) / 10000000ull);
168 
169     if (psutil_GetTickCount64 != NULL) {
170         // Windows >= Vista
171         uptime = psutil_GetTickCount64() / 1000ull;
172     }
173     else {
174         // Windows XP.
175         // GetTickCount() time will wrap around to zero if the
176         // system is run continuously for 49.7 days.
177         uptime = (ULONGLONG)GetTickCount() / 1000ull;
178     }
179     return Py_BuildValue("K", pt - uptime);
180 }
181 
182 
183 /*
184  * Return 1 if PID exists in the current process list, else 0.
185  */
186 static PyObject *
psutil_pid_exists(PyObject * self,PyObject * args)187 psutil_pid_exists(PyObject *self, PyObject *args) {
188     long pid;
189     int status;
190 
191     if (! PyArg_ParseTuple(args, "l", &pid))
192         return NULL;
193 
194     status = psutil_pid_is_running(pid);
195     if (-1 == status)
196         return NULL; // exception raised in psutil_pid_is_running()
197     return PyBool_FromLong(status);
198 }
199 
200 
201 /*
202  * Return a Python list of all the PIDs running on the system.
203  */
204 static PyObject *
psutil_pids(PyObject * self,PyObject * args)205 psutil_pids(PyObject *self, PyObject *args) {
206     DWORD *proclist = NULL;
207     DWORD numberOfReturnedPIDs;
208     DWORD i;
209     PyObject *py_pid = NULL;
210     PyObject *py_retlist = PyList_New(0);
211 
212     if (py_retlist == NULL)
213         return NULL;
214     proclist = psutil_get_pids(&numberOfReturnedPIDs);
215     if (proclist == NULL)
216         goto error;
217 
218     for (i = 0; i < numberOfReturnedPIDs; i++) {
219         py_pid = Py_BuildValue("I", proclist[i]);
220         if (!py_pid)
221             goto error;
222         if (PyList_Append(py_retlist, py_pid))
223             goto error;
224         Py_CLEAR(py_pid);
225     }
226 
227     // free C array allocated for PIDs
228     free(proclist);
229     return py_retlist;
230 
231 error:
232     Py_XDECREF(py_pid);
233     Py_DECREF(py_retlist);
234     if (proclist != NULL)
235         free(proclist);
236     return NULL;
237 }
238 
239 
240 /*
241  * Kill a process given its PID.
242  */
243 static PyObject *
psutil_proc_kill(PyObject * self,PyObject * args)244 psutil_proc_kill(PyObject *self, PyObject *args) {
245     HANDLE hProcess;
246     long pid;
247 
248     if (! PyArg_ParseTuple(args, "l", &pid))
249         return NULL;
250     if (pid == 0)
251         return AccessDenied("");
252 
253     hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
254     if (hProcess == NULL) {
255         if (GetLastError() == ERROR_INVALID_PARAMETER) {
256             // see https://github.com/giampaolo/psutil/issues/24
257             psutil_debug("OpenProcess -> ERROR_INVALID_PARAMETER turned "
258                          "into NoSuchProcess");
259             NoSuchProcess("");
260         }
261         else {
262             PyErr_SetFromWindowsErr(0);
263         }
264         return NULL;
265     }
266 
267     if (! TerminateProcess(hProcess, SIGTERM)) {
268         // ERROR_ACCESS_DENIED may happen if the process already died. See:
269         // https://github.com/giampaolo/psutil/issues/1099
270         if (GetLastError() != ERROR_ACCESS_DENIED) {
271             PyErr_SetFromOSErrnoWithSyscall("TerminateProcess");
272             return NULL;
273         }
274     }
275 
276     CloseHandle(hProcess);
277     Py_RETURN_NONE;
278 }
279 
280 
281 /*
282  * Wait for process to terminate and return its exit code.
283  */
284 static PyObject *
psutil_proc_wait(PyObject * self,PyObject * args)285 psutil_proc_wait(PyObject *self, PyObject *args) {
286     HANDLE hProcess;
287     DWORD ExitCode;
288     DWORD retVal;
289     long pid;
290     long timeout;
291 
292     if (! PyArg_ParseTuple(args, "ll", &pid, &timeout))
293         return NULL;
294     if (pid == 0)
295         return AccessDenied("");
296 
297     hProcess = OpenProcess(SYNCHRONIZE | PROCESS_QUERY_INFORMATION,
298                            FALSE, pid);
299     if (hProcess == NULL) {
300         if (GetLastError() == ERROR_INVALID_PARAMETER) {
301             // no such process; we do not want to raise NSP but
302             // return None instead.
303             Py_RETURN_NONE;
304         }
305         else
306             return PyErr_SetFromWindowsErr(0);
307     }
308 
309     // wait until the process has terminated
310     Py_BEGIN_ALLOW_THREADS
311     retVal = WaitForSingleObject(hProcess, timeout);
312     Py_END_ALLOW_THREADS
313 
314     // handle return code
315     if (retVal == WAIT_FAILED) {
316         PyErr_SetFromOSErrnoWithSyscall("WaitForSingleObject");
317         CloseHandle(hProcess);
318         return NULL;
319     }
320     if (retVal == WAIT_TIMEOUT) {
321         PyErr_SetString(TimeoutExpired,
322                         "WaitForSingleObject() returned WAIT_TIMEOUT");
323         CloseHandle(hProcess);
324         return NULL;
325     }
326     if (retVal == WAIT_ABANDONED) {
327         psutil_debug("WaitForSingleObject() -> WAIT_ABANDONED");
328         PyErr_SetString(TimeoutAbandoned,
329                         "WaitForSingleObject() returned WAIT_ABANDONED");
330         CloseHandle(hProcess);
331         return NULL;
332     }
333 
334     // WaitForSingleObject() returned WAIT_OBJECT_0. It means the
335     // process is gone so we can get its process exit code. The PID
336     // may still stick around though but we'll handle that from Python.
337     if (GetExitCodeProcess(hProcess, &ExitCode) == 0) {
338         PyErr_SetFromOSErrnoWithSyscall("GetExitCodeProcess");
339         CloseHandle(hProcess);
340         return NULL;
341     }
342 
343     CloseHandle(hProcess);
344 
345 #if PY_MAJOR_VERSION >= 3
346     return PyLong_FromLong((long) ExitCode);
347 #else
348     return PyInt_FromLong((long) ExitCode);
349 #endif
350 }
351 
352 
353 /*
354  * Return a Python tuple (user_time, kernel_time)
355  */
356 static PyObject *
psutil_proc_cpu_times(PyObject * self,PyObject * args)357 psutil_proc_cpu_times(PyObject *self, PyObject *args) {
358     long        pid;
359     HANDLE      hProcess;
360     FILETIME    ftCreate, ftExit, ftKernel, ftUser;
361 
362     if (! PyArg_ParseTuple(args, "l", &pid))
363         return NULL;
364 
365     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
366 
367     if (hProcess == NULL)
368         return NULL;
369     if (! GetProcessTimes(hProcess, &ftCreate, &ftExit, &ftKernel, &ftUser)) {
370         if (GetLastError() == ERROR_ACCESS_DENIED) {
371             // usually means the process has died so we throw a NoSuchProcess
372             // here
373             NoSuchProcess("");
374         }
375         else {
376             PyErr_SetFromWindowsErr(0);
377         }
378         CloseHandle(hProcess);
379         return NULL;
380     }
381 
382     CloseHandle(hProcess);
383 
384     /*
385      * User and kernel times are represented as a FILETIME structure
386      * which contains a 64-bit value representing the number of
387      * 100-nanosecond intervals since January 1, 1601 (UTC):
388      * http://msdn.microsoft.com/en-us/library/ms724284(VS.85).aspx
389      * To convert it into a float representing the seconds that the
390      * process has executed in user/kernel mode I borrowed the code
391      * below from Python's Modules/posixmodule.c
392      */
393     return Py_BuildValue(
394        "(dd)",
395        (double)(ftUser.dwHighDateTime * 429.4967296 + \
396                 ftUser.dwLowDateTime * 1e-7),
397        (double)(ftKernel.dwHighDateTime * 429.4967296 + \
398                 ftKernel.dwLowDateTime * 1e-7)
399    );
400 }
401 
402 
403 /*
404  * Return a Python float indicating the process create time expressed in
405  * seconds since the epoch.
406  */
407 static PyObject *
psutil_proc_create_time(PyObject * self,PyObject * args)408 psutil_proc_create_time(PyObject *self, PyObject *args) {
409     long        pid;
410     long long   unix_time;
411     HANDLE      hProcess;
412     FILETIME    ftCreate, ftExit, ftKernel, ftUser;
413 
414     if (! PyArg_ParseTuple(args, "l", &pid))
415         return NULL;
416 
417     // special case for PIDs 0 and 4, return system boot time
418     if (0 == pid || 4 == pid)
419         return psutil_boot_time(NULL, NULL);
420 
421     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
422     if (hProcess == NULL)
423         return NULL;
424     if (! GetProcessTimes(hProcess, &ftCreate, &ftExit, &ftKernel, &ftUser)) {
425         if (GetLastError() == ERROR_ACCESS_DENIED) {
426             // usually means the process has died so we throw a
427             // NoSuchProcess here
428             NoSuchProcess("");
429         }
430         else {
431             PyErr_SetFromWindowsErr(0);
432         }
433         CloseHandle(hProcess);
434         return NULL;
435     }
436 
437     CloseHandle(hProcess);
438 
439     /*
440     // Make sure the process is not gone as OpenProcess alone seems to be
441     // unreliable in doing so (it seems a previous call to p.wait() makes
442     // it unreliable).
443     // This check is important as creation time is used to make sure the
444     // process is still running.
445     ret = GetExitCodeProcess(hProcess, &exitCode);
446     CloseHandle(hProcess);
447     if (ret != 0) {
448         if (exitCode != STILL_ACTIVE)
449             return NoSuchProcess("");
450     }
451     else {
452         // Ignore access denied as it means the process is still alive.
453         // For all other errors, we want an exception.
454         if (GetLastError() != ERROR_ACCESS_DENIED)
455             return PyErr_SetFromWindowsErr(0);
456     }
457     */
458 
459     // Convert the FILETIME structure to a Unix time.
460     // It's the best I could find by googling and borrowing code here
461     // and there. The time returned has a precision of 1 second.
462     unix_time = ((LONGLONG)ftCreate.dwHighDateTime) << 32;
463     unix_time += ftCreate.dwLowDateTime - 116444736000000000LL;
464     unix_time /= 10000000;
465     return Py_BuildValue("d", (double)unix_time);
466 }
467 
468 
469 /*
470  * Return the number of active, logical CPUs.
471  */
472 static PyObject *
psutil_cpu_count_logical(PyObject * self,PyObject * args)473 psutil_cpu_count_logical(PyObject *self, PyObject *args) {
474     unsigned int ncpus;
475 
476     ncpus = psutil_get_num_cpus(0);
477     if (ncpus != 0)
478         return Py_BuildValue("I", ncpus);
479     else
480         Py_RETURN_NONE;  // mimick os.cpu_count()
481 }
482 
483 
484 /*
485  * Return the number of physical CPU cores (hyper-thread CPUs count
486  * is excluded).
487  */
488 static PyObject *
psutil_cpu_count_phys(PyObject * self,PyObject * args)489 psutil_cpu_count_phys(PyObject *self, PyObject *args) {
490     DWORD rc;
491     PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX buffer = NULL;
492     PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX ptr = NULL;
493     DWORD length = 0;
494     DWORD offset = 0;
495     DWORD ncpus = 0;
496     DWORD prev_processor_info_size = 0;
497 
498     // GetLogicalProcessorInformationEx() is available from Windows 7
499     // onward. Differently from GetLogicalProcessorInformation()
500     // it supports process groups, meaning this is able to report more
501     // than 64 CPUs. See:
502     // https://bugs.python.org/issue33166
503     if (psutil_GetLogicalProcessorInformationEx == NULL) {
504         psutil_debug("Win < 7; cpu_count_phys() forced to None");
505         Py_RETURN_NONE;
506     }
507 
508     while (1) {
509         rc = psutil_GetLogicalProcessorInformationEx(
510             RelationAll, buffer, &length);
511         if (rc == FALSE) {
512             if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
513                 if (buffer) {
514                     free(buffer);
515                 }
516                 buffer = \
517                     (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)malloc(length);
518                 if (NULL == buffer) {
519                     PyErr_NoMemory();
520                     return NULL;
521                 }
522             }
523             else {
524                 psutil_debug("GetLogicalProcessorInformationEx() returned ",
525                              GetLastError());
526                 goto return_none;
527             }
528         }
529         else {
530             break;
531         }
532     }
533 
534     ptr = buffer;
535     while (offset < length) {
536         // Advance ptr by the size of the previous
537         // SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX struct.
538         ptr = (SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX*)\
539             (((char*)ptr) + prev_processor_info_size);
540 
541         if (ptr->Relationship == RelationProcessorCore) {
542             ncpus += 1;
543         }
544 
545         // When offset == length, we've reached the last processor
546         // info struct in the buffer.
547         offset += ptr->Size;
548         prev_processor_info_size = ptr->Size;
549     }
550 
551     free(buffer);
552     if (ncpus != 0) {
553         return Py_BuildValue("I", ncpus);
554     }
555     else {
556         psutil_debug("GetLogicalProcessorInformationEx() count was 0");
557         Py_RETURN_NONE;  // mimick os.cpu_count()
558     }
559 
560 return_none:
561     if (buffer != NULL)
562         free(buffer);
563     Py_RETURN_NONE;
564 }
565 
566 
567 /*
568  * Return process cmdline as a Python list of cmdline arguments.
569  */
570 static PyObject *
psutil_proc_cmdline(PyObject * self,PyObject * args,PyObject * kwdict)571 psutil_proc_cmdline(PyObject *self, PyObject *args, PyObject *kwdict) {
572     long pid;
573     int pid_return;
574     int use_peb;
575     PyObject *py_usepeb = Py_True;
576     static char *keywords[] = {"pid", "use_peb", NULL};
577 
578     if (!PyArg_ParseTupleAndKeywords(args, kwdict, "i|O",
579                                      keywords, &pid, &py_usepeb)) {
580         return NULL;
581     }
582     if ((pid == 0) || (pid == 4))
583         return Py_BuildValue("[]");
584 
585     pid_return = psutil_pid_is_running(pid);
586     if (pid_return == 0)
587         return NoSuchProcess("");
588     if (pid_return == -1)
589         return NULL;
590 
591     use_peb = (py_usepeb == Py_True) ? 1 : 0;
592     return psutil_get_cmdline(pid, use_peb);
593 }
594 
595 
596 /*
597  * Return process cmdline as a Python list of cmdline arguments.
598  */
599 static PyObject *
psutil_proc_environ(PyObject * self,PyObject * args)600 psutil_proc_environ(PyObject *self, PyObject *args) {
601     long pid;
602     int pid_return;
603 
604     if (! PyArg_ParseTuple(args, "l", &pid))
605         return NULL;
606     if ((pid == 0) || (pid == 4))
607         return Py_BuildValue("s", "");
608 
609     pid_return = psutil_pid_is_running(pid);
610     if (pid_return == 0)
611         return NoSuchProcess("");
612     if (pid_return == -1)
613         return NULL;
614 
615     return psutil_get_environ(pid);
616 }
617 
618 
619 /*
620  * Return process executable path.
621  */
622 static PyObject *
psutil_proc_exe(PyObject * self,PyObject * args)623 psutil_proc_exe(PyObject *self, PyObject *args) {
624     long pid;
625     HANDLE hProcess;
626     wchar_t exe[MAX_PATH];
627 #if (_WIN32_WINNT >= 0x0600)  // >= Vista
628     unsigned int size = sizeof(exe);
629 #endif
630 
631     if (! PyArg_ParseTuple(args, "l", &pid))
632         return NULL;
633     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
634     if (NULL == hProcess)
635         return NULL;
636 
637     // Here we differentiate between XP and Vista+ because
638     // QueryFullProcessImageNameW is better than GetProcessImageFileNameW
639     // (avoid using QueryDosDevice on the returned path), see:
640     // https://github.com/giampaolo/psutil/issues/1394
641 #if (_WIN32_WINNT >= 0x0600)  // Windows >= Vista
642     memset(exe, 0, MAX_PATH);
643     if (QueryFullProcessImageNameW(hProcess, 0, exe, &size) == 0) {
644         PyErr_SetFromOSErrnoWithSyscall("QueryFullProcessImageNameW");
645         CloseHandle(hProcess);
646         return NULL;
647     }
648 #else  // Windows XP
649     if (GetProcessImageFileNameW(hProcess, exe, MAX_PATH) == 0) {
650         // see: https://github.com/giampaolo/psutil/issues/1394
651         if (GetLastError() == 0)
652             PyErr_SetFromWindowsErr(ERROR_ACCESS_DENIED);
653         else
654             PyErr_SetFromOSErrnoWithSyscall("GetProcessImageFileNameW");
655         CloseHandle(hProcess);
656         return NULL;
657     }
658 #endif
659     CloseHandle(hProcess);
660     return PyUnicode_FromWideChar(exe, wcslen(exe));
661 }
662 
663 
664 /*
665  * Return process base name.
666  * Note: psutil_proc_exe() is attempted first because it's faster
667  * but it raise AccessDenied for processes owned by other users
668  * in which case we fall back on using this.
669  */
670 static PyObject *
psutil_proc_name(PyObject * self,PyObject * args)671 psutil_proc_name(PyObject *self, PyObject *args) {
672     long pid;
673     int ok;
674     PROCESSENTRY32W pentry;
675     HANDLE hSnapShot;
676 
677     if (! PyArg_ParseTuple(args, "l", &pid))
678         return NULL;
679     hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, pid);
680     if (hSnapShot == INVALID_HANDLE_VALUE)
681         return PyErr_SetFromOSErrnoWithSyscall("CreateToolhelp32Snapshot");
682     pentry.dwSize = sizeof(PROCESSENTRY32W);
683     ok = Process32FirstW(hSnapShot, &pentry);
684     if (! ok) {
685         PyErr_SetFromOSErrnoWithSyscall("Process32FirstW");
686         CloseHandle(hSnapShot);
687         return NULL;
688     }
689     while (ok) {
690         if (pentry.th32ProcessID == pid) {
691             CloseHandle(hSnapShot);
692             return PyUnicode_FromWideChar(
693                 pentry.szExeFile, wcslen(pentry.szExeFile));
694         }
695         ok = Process32NextW(hSnapShot, &pentry);
696     }
697 
698     CloseHandle(hSnapShot);
699     NoSuchProcess("");
700     return NULL;
701 }
702 
703 
704 /*
705  * Return process memory information as a Python tuple.
706  */
707 static PyObject *
psutil_proc_memory_info(PyObject * self,PyObject * args)708 psutil_proc_memory_info(PyObject *self, PyObject *args) {
709     HANDLE hProcess;
710     DWORD pid;
711 #if (_WIN32_WINNT >= 0x0501)  // Windows XP with SP2
712     PROCESS_MEMORY_COUNTERS_EX cnt;
713 #else
714     PROCESS_MEMORY_COUNTERS cnt;
715 #endif
716     SIZE_T private = 0;
717 
718     if (! PyArg_ParseTuple(args, "l", &pid))
719         return NULL;
720 
721     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
722     if (NULL == hProcess)
723         return NULL;
724 
725     if (! GetProcessMemoryInfo(hProcess, (PPROCESS_MEMORY_COUNTERS)&cnt,
726                                sizeof(cnt))) {
727         PyErr_SetFromWindowsErr(0);
728         CloseHandle(hProcess);
729         return NULL;
730     }
731 
732 #if (_WIN32_WINNT >= 0x0501)  // Windows XP with SP2
733     private = cnt.PrivateUsage;
734 #endif
735 
736     CloseHandle(hProcess);
737 
738     // PROCESS_MEMORY_COUNTERS values are defined as SIZE_T which on 64bits
739     // is an (unsigned long long) and on 32bits is an (unsigned int).
740     // "_WIN64" is defined if we're running a 64bit Python interpreter not
741     // exclusively if the *system* is 64bit.
742 #if defined(_WIN64)
743     return Py_BuildValue(
744         "(kKKKKKKKKK)",
745         cnt.PageFaultCount,  // unsigned long
746         (unsigned long long)cnt.PeakWorkingSetSize,
747         (unsigned long long)cnt.WorkingSetSize,
748         (unsigned long long)cnt.QuotaPeakPagedPoolUsage,
749         (unsigned long long)cnt.QuotaPagedPoolUsage,
750         (unsigned long long)cnt.QuotaPeakNonPagedPoolUsage,
751         (unsigned long long)cnt.QuotaNonPagedPoolUsage,
752         (unsigned long long)cnt.PagefileUsage,
753         (unsigned long long)cnt.PeakPagefileUsage,
754         (unsigned long long)private);
755 #else
756     return Py_BuildValue(
757         "(kIIIIIIIII)",
758         cnt.PageFaultCount,    // unsigned long
759         (unsigned int)cnt.PeakWorkingSetSize,
760         (unsigned int)cnt.WorkingSetSize,
761         (unsigned int)cnt.QuotaPeakPagedPoolUsage,
762         (unsigned int)cnt.QuotaPagedPoolUsage,
763         (unsigned int)cnt.QuotaPeakNonPagedPoolUsage,
764         (unsigned int)cnt.QuotaNonPagedPoolUsage,
765         (unsigned int)cnt.PagefileUsage,
766         (unsigned int)cnt.PeakPagefileUsage,
767         (unsigned int)private);
768 #endif
769 }
770 
771 
772 static int
psutil_GetProcWsetInformation(DWORD pid,HANDLE hProcess,PMEMORY_WORKING_SET_INFORMATION * wSetInfo)773 psutil_GetProcWsetInformation(
774         DWORD pid,
775         HANDLE hProcess,
776         PMEMORY_WORKING_SET_INFORMATION *wSetInfo)
777 {
778     NTSTATUS status;
779     PVOID buffer;
780     SIZE_T bufferSize;
781 
782     bufferSize = 0x8000;
783     buffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, bufferSize);
784 
785     while ((status = psutil_NtQueryVirtualMemory(
786             hProcess,
787             NULL,
788             MemoryWorkingSetInformation,
789             buffer,
790             bufferSize,
791             NULL)) == STATUS_INFO_LENGTH_MISMATCH)
792     {
793         HeapFree(GetProcessHeap(), 0, buffer);
794         bufferSize *= 2;
795         psutil_debug("NtQueryVirtualMemory increase bufsize %zd", bufferSize);
796         // Fail if we're resizing the buffer to something very large.
797         if (bufferSize > 256 * 1024 * 1024) {
798             PyErr_SetString(PyExc_RuntimeError,
799                             "NtQueryVirtualMemory bufsize is too large");
800             return 1;
801         }
802         buffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, bufferSize);
803     }
804 
805     if (!NT_SUCCESS(status)) {
806         if (status == STATUS_ACCESS_DENIED) {
807             AccessDenied("originated from NtQueryVirtualMemory");
808         }
809         else if (psutil_pid_is_running(pid) == 0) {
810             NoSuchProcess("");
811         }
812         else {
813             PyErr_Clear();
814             psutil_SetFromNTStatusErr(
815                 status, "NtQueryVirtualMemory(MemoryWorkingSetInformation)");
816         }
817         HeapFree(GetProcessHeap(), 0, buffer);
818         return 1;
819     }
820 
821     *wSetInfo = (PMEMORY_WORKING_SET_INFORMATION)buffer;
822     return 0;
823 }
824 
825 
826 /*
827  * Returns the USS of the process.
828  * Reference:
829  * https://dxr.mozilla.org/mozilla-central/source/xpcom/base/
830  *     nsMemoryReporterManager.cpp
831  */
832 static PyObject *
psutil_proc_memory_uss(PyObject * self,PyObject * args)833 psutil_proc_memory_uss(PyObject *self, PyObject *args) {
834     DWORD pid;
835     HANDLE hProcess;
836     PSUTIL_PROCESS_WS_COUNTERS wsCounters;
837     PMEMORY_WORKING_SET_INFORMATION wsInfo;
838     ULONG_PTR i;
839 
840     if (! PyArg_ParseTuple(args, "l", &pid))
841         return NULL;
842     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
843     if (hProcess == NULL)
844         return NULL;
845 
846     if (psutil_GetProcWsetInformation(pid, hProcess, &wsInfo) != 0) {
847         CloseHandle(hProcess);
848         return NULL;
849     }
850     memset(&wsCounters, 0, sizeof(PSUTIL_PROCESS_WS_COUNTERS));
851 
852     for (i = 0; i < wsInfo->NumberOfEntries; i++) {
853         // This is what ProcessHacker does.
854         /*
855         wsCounters.NumberOfPages++;
856         if (wsInfo->WorkingSetInfo[i].ShareCount > 1)
857             wsCounters.NumberOfSharedPages++;
858         if (wsInfo->WorkingSetInfo[i].ShareCount == 0)
859             wsCounters.NumberOfPrivatePages++;
860         if (wsInfo->WorkingSetInfo[i].Shared)
861             wsCounters.NumberOfShareablePages++;
862         */
863 
864         // This is what we do: count shared pages that only one process
865         // is using as private (USS).
866         if (!wsInfo->WorkingSetInfo[i].Shared ||
867                 wsInfo->WorkingSetInfo[i].ShareCount <= 1) {
868             wsCounters.NumberOfPrivatePages++;
869         }
870     }
871 
872     HeapFree(GetProcessHeap(), 0, wsInfo);
873     CloseHandle(hProcess);
874 
875     return Py_BuildValue("I", wsCounters.NumberOfPrivatePages);
876 }
877 
878 
879 /*
880  * Return a Python integer indicating the total amount of physical memory
881  * in bytes.
882  */
883 static PyObject *
psutil_virtual_mem(PyObject * self,PyObject * args)884 psutil_virtual_mem(PyObject *self, PyObject *args) {
885     MEMORYSTATUSEX memInfo;
886     memInfo.dwLength = sizeof(MEMORYSTATUSEX);
887 
888     if (! GlobalMemoryStatusEx(&memInfo))
889         return PyErr_SetFromWindowsErr(0);
890     return Py_BuildValue("(LLLLLL)",
891                          memInfo.ullTotalPhys,      // total
892                          memInfo.ullAvailPhys,      // avail
893                          memInfo.ullTotalPageFile,  // total page file
894                          memInfo.ullAvailPageFile,  // avail page file
895                          memInfo.ullTotalVirtual,   // total virtual
896                          memInfo.ullAvailVirtual);  // avail virtual
897 }
898 
899 /*
900  * Retrieves system CPU timing information as a (user, system, idle)
901  * tuple. On a multiprocessor system, the values returned are the
902  * sum of the designated times across all processors.
903  */
904 static PyObject *
psutil_cpu_times(PyObject * self,PyObject * args)905 psutil_cpu_times(PyObject *self, PyObject *args) {
906     double idle, kernel, user, system;
907     FILETIME idle_time, kernel_time, user_time;
908 
909     if (!GetSystemTimes(&idle_time, &kernel_time, &user_time))
910         return PyErr_SetFromWindowsErr(0);
911 
912     idle = (double)((HI_T * idle_time.dwHighDateTime) + \
913                    (LO_T * idle_time.dwLowDateTime));
914     user = (double)((HI_T * user_time.dwHighDateTime) + \
915                    (LO_T * user_time.dwLowDateTime));
916     kernel = (double)((HI_T * kernel_time.dwHighDateTime) + \
917                      (LO_T * kernel_time.dwLowDateTime));
918 
919     // Kernel time includes idle time.
920     // We return only busy kernel time subtracting idle time from
921     // kernel time.
922     system = (kernel - idle);
923     return Py_BuildValue("(ddd)", user, system, idle);
924 }
925 
926 
927 /*
928  * Same as above but for all system CPUs.
929  */
930 static PyObject *
psutil_per_cpu_times(PyObject * self,PyObject * args)931 psutil_per_cpu_times(PyObject *self, PyObject *args) {
932     double idle, kernel, systemt, user, interrupt, dpc;
933     NTSTATUS status;
934     _SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *sppi = NULL;
935     UINT i;
936     unsigned int ncpus;
937     PyObject *py_tuple = NULL;
938     PyObject *py_retlist = PyList_New(0);
939 
940     if (py_retlist == NULL)
941         return NULL;
942 
943     // retrieves number of processors
944     ncpus = psutil_get_num_cpus(1);
945     if (ncpus == 0)
946         goto error;
947 
948     // allocates an array of _SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION
949     // structures, one per processor
950     sppi = (_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *) \
951         malloc(ncpus * sizeof(_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION));
952     if (sppi == NULL) {
953         PyErr_NoMemory();
954         goto error;
955     }
956 
957     // gets cpu time informations
958     status = psutil_NtQuerySystemInformation(
959         SystemProcessorPerformanceInformation,
960         sppi,
961         ncpus * sizeof(_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION),
962         NULL);
963     if (! NT_SUCCESS(status)) {
964         psutil_SetFromNTStatusErr(
965             status,
966             "NtQuerySystemInformation(SystemProcessorPerformanceInformation)"
967         );
968         goto error;
969     }
970 
971     // computes system global times summing each
972     // processor value
973     idle = user = kernel = interrupt = dpc = 0;
974     for (i = 0; i < ncpus; i++) {
975         py_tuple = NULL;
976         user = (double)((HI_T * sppi[i].UserTime.HighPart) +
977                        (LO_T * sppi[i].UserTime.LowPart));
978         idle = (double)((HI_T * sppi[i].IdleTime.HighPart) +
979                        (LO_T * sppi[i].IdleTime.LowPart));
980         kernel = (double)((HI_T * sppi[i].KernelTime.HighPart) +
981                          (LO_T * sppi[i].KernelTime.LowPart));
982         interrupt = (double)((HI_T * sppi[i].InterruptTime.HighPart) +
983                             (LO_T * sppi[i].InterruptTime.LowPart));
984         dpc = (double)((HI_T * sppi[i].DpcTime.HighPart) +
985                       (LO_T * sppi[i].DpcTime.LowPart));
986 
987         // kernel time includes idle time on windows
988         // we return only busy kernel time subtracting
989         // idle time from kernel time
990         systemt = kernel - idle;
991         py_tuple = Py_BuildValue(
992             "(ddddd)",
993             user,
994             systemt,
995             idle,
996             interrupt,
997             dpc
998         );
999         if (!py_tuple)
1000             goto error;
1001         if (PyList_Append(py_retlist, py_tuple))
1002             goto error;
1003         Py_CLEAR(py_tuple);
1004     }
1005 
1006     free(sppi);
1007     return py_retlist;
1008 
1009 error:
1010     Py_XDECREF(py_tuple);
1011     Py_DECREF(py_retlist);
1012     if (sppi)
1013         free(sppi);
1014     return NULL;
1015 }
1016 
1017 
1018 /*
1019  * Return process current working directory as a Python string.
1020  */
1021 static PyObject *
psutil_proc_cwd(PyObject * self,PyObject * args)1022 psutil_proc_cwd(PyObject *self, PyObject *args) {
1023     long pid;
1024     int pid_return;
1025 
1026     if (! PyArg_ParseTuple(args, "l", &pid))
1027         return NULL;
1028 
1029     pid_return = psutil_pid_is_running(pid);
1030     if (pid_return == 0)
1031         return NoSuchProcess("");
1032     if (pid_return == -1)
1033         return NULL;
1034 
1035     return psutil_get_cwd(pid);
1036 }
1037 
1038 
1039 /*
1040  * Resume or suspends a process
1041  */
1042 static PyObject *
psutil_proc_suspend_or_resume(PyObject * self,PyObject * args)1043 psutil_proc_suspend_or_resume(PyObject *self, PyObject *args) {
1044     long pid;
1045     NTSTATUS status;
1046     HANDLE hProcess;
1047     PyObject* suspend;
1048 
1049     if (! PyArg_ParseTuple(args, "lO", &pid, &suspend))
1050         return NULL;
1051 
1052     hProcess = psutil_handle_from_pid(pid, PROCESS_SUSPEND_RESUME);
1053     if (hProcess == NULL)
1054         return NULL;
1055 
1056     if (PyObject_IsTrue(suspend))
1057         status = psutil_NtSuspendProcess(hProcess);
1058     else
1059         status = psutil_NtResumeProcess(hProcess);
1060 
1061     if (! NT_SUCCESS(status)) {
1062         CloseHandle(hProcess);
1063         return psutil_SetFromNTStatusErr(status, "NtSuspend|ResumeProcess");
1064     }
1065 
1066     CloseHandle(hProcess);
1067     Py_RETURN_NONE;
1068 }
1069 
1070 
1071 static PyObject *
psutil_proc_threads(PyObject * self,PyObject * args)1072 psutil_proc_threads(PyObject *self, PyObject *args) {
1073     HANDLE hThread;
1074     THREADENTRY32 te32 = {0};
1075     long pid;
1076     int pid_return;
1077     int rc;
1078     FILETIME ftDummy, ftKernel, ftUser;
1079     HANDLE hThreadSnap = NULL;
1080     PyObject *py_tuple = NULL;
1081     PyObject *py_retlist = PyList_New(0);
1082 
1083     if (py_retlist == NULL)
1084         return NULL;
1085     if (! PyArg_ParseTuple(args, "l", &pid))
1086         goto error;
1087     if (pid == 0) {
1088         // raise AD instead of returning 0 as procexp is able to
1089         // retrieve useful information somehow
1090         AccessDenied("");
1091         goto error;
1092     }
1093 
1094     pid_return = psutil_pid_is_running(pid);
1095     if (pid_return == 0) {
1096         NoSuchProcess("");
1097         goto error;
1098     }
1099     if (pid_return == -1)
1100         goto error;
1101 
1102     hThreadSnap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
1103     if (hThreadSnap == INVALID_HANDLE_VALUE) {
1104         PyErr_SetFromOSErrnoWithSyscall("CreateToolhelp32Snapshot");
1105         goto error;
1106     }
1107 
1108     // Fill in the size of the structure before using it
1109     te32.dwSize = sizeof(THREADENTRY32);
1110 
1111     if (! Thread32First(hThreadSnap, &te32)) {
1112         PyErr_SetFromOSErrnoWithSyscall("Thread32First");
1113         goto error;
1114     }
1115 
1116     // Walk the thread snapshot to find all threads of the process.
1117     // If the thread belongs to the process, increase the counter.
1118     do {
1119         if (te32.th32OwnerProcessID == pid) {
1120             py_tuple = NULL;
1121             hThread = NULL;
1122             hThread = OpenThread(THREAD_QUERY_INFORMATION,
1123                                  FALSE, te32.th32ThreadID);
1124             if (hThread == NULL) {
1125                 // thread has disappeared on us
1126                 continue;
1127             }
1128 
1129             rc = GetThreadTimes(hThread, &ftDummy, &ftDummy, &ftKernel,
1130                                 &ftUser);
1131             if (rc == 0) {
1132                 PyErr_SetFromOSErrnoWithSyscall("GetThreadTimes");
1133                 goto error;
1134             }
1135 
1136             /*
1137              * User and kernel times are represented as a FILETIME structure
1138              * which contains a 64-bit value representing the number of
1139              * 100-nanosecond intervals since January 1, 1601 (UTC):
1140              * http://msdn.microsoft.com/en-us/library/ms724284(VS.85).aspx
1141              * To convert it into a float representing the seconds that the
1142              * process has executed in user/kernel mode I borrowed the code
1143              * below from Python's Modules/posixmodule.c
1144              */
1145             py_tuple = Py_BuildValue(
1146                 "kdd",
1147                 te32.th32ThreadID,
1148                 (double)(ftUser.dwHighDateTime * 429.4967296 + \
1149                          ftUser.dwLowDateTime * 1e-7),
1150                 (double)(ftKernel.dwHighDateTime * 429.4967296 + \
1151                          ftKernel.dwLowDateTime * 1e-7));
1152             if (!py_tuple)
1153                 goto error;
1154             if (PyList_Append(py_retlist, py_tuple))
1155                 goto error;
1156             Py_CLEAR(py_tuple);
1157 
1158             CloseHandle(hThread);
1159         }
1160     } while (Thread32Next(hThreadSnap, &te32));
1161 
1162     CloseHandle(hThreadSnap);
1163     return py_retlist;
1164 
1165 error:
1166     Py_XDECREF(py_tuple);
1167     Py_DECREF(py_retlist);
1168     if (hThread != NULL)
1169         CloseHandle(hThread);
1170     if (hThreadSnap != NULL)
1171         CloseHandle(hThreadSnap);
1172     return NULL;
1173 }
1174 
1175 
1176 static PyObject *
psutil_proc_open_files(PyObject * self,PyObject * args)1177 psutil_proc_open_files(PyObject *self, PyObject *args) {
1178     long       pid;
1179     HANDLE     processHandle;
1180     DWORD      access = PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION;
1181     PyObject  *py_retlist;
1182 
1183     if (! PyArg_ParseTuple(args, "l", &pid))
1184         return NULL;
1185 
1186     processHandle = psutil_handle_from_pid(pid, access);
1187     if (processHandle == NULL)
1188         return NULL;
1189 
1190     py_retlist = psutil_get_open_files(pid, processHandle);
1191     if (py_retlist == NULL) {
1192         PyErr_SetFromWindowsErr(0);
1193         CloseHandle(processHandle);
1194         return NULL;
1195     }
1196 
1197     CloseHandle(processHandle);
1198     return py_retlist;
1199 }
1200 
1201 
1202 /*
1203  Accept a filename's drive in native  format like "\Device\HarddiskVolume1\"
1204  and return the corresponding drive letter (e.g. "C:\\").
1205  If no match is found return an empty string.
1206 */
1207 static PyObject *
psutil_win32_QueryDosDevice(PyObject * self,PyObject * args)1208 psutil_win32_QueryDosDevice(PyObject *self, PyObject *args) {
1209     LPCTSTR   lpDevicePath;
1210     TCHAR d = TEXT('A');
1211     TCHAR     szBuff[5];
1212 
1213     if (!PyArg_ParseTuple(args, "s", &lpDevicePath))
1214         return NULL;
1215 
1216     while (d <= TEXT('Z')) {
1217         TCHAR szDeviceName[3] = {d, TEXT(':'), TEXT('\0')};
1218         TCHAR szTarget[512] = {0};
1219         if (QueryDosDevice(szDeviceName, szTarget, 511) != 0) {
1220             if (_tcscmp(lpDevicePath, szTarget) == 0) {
1221                 _stprintf_s(szBuff, _countof(szBuff), TEXT("%c:"), d);
1222                 return Py_BuildValue("s", szBuff);
1223             }
1224         }
1225         d++;
1226     }
1227     return Py_BuildValue("s", "");
1228 }
1229 
1230 
1231 /*
1232  * Return process username as a "DOMAIN//USERNAME" string.
1233  */
1234 static PyObject *
psutil_proc_username(PyObject * self,PyObject * args)1235 psutil_proc_username(PyObject *self, PyObject *args) {
1236     long pid;
1237     HANDLE processHandle = NULL;
1238     HANDLE tokenHandle = NULL;
1239     PTOKEN_USER user = NULL;
1240     ULONG bufferSize;
1241     WCHAR *name = NULL;
1242     WCHAR *domainName = NULL;
1243     ULONG nameSize;
1244     ULONG domainNameSize;
1245     SID_NAME_USE nameUse;
1246     PyObject *py_username = NULL;
1247     PyObject *py_domain = NULL;
1248     PyObject *py_tuple = NULL;
1249 
1250     if (! PyArg_ParseTuple(args, "l", &pid))
1251         return NULL;
1252 
1253     processHandle = psutil_handle_from_pid(
1254         pid, PROCESS_QUERY_LIMITED_INFORMATION);
1255     if (processHandle == NULL)
1256         return NULL;
1257 
1258     if (!OpenProcessToken(processHandle, TOKEN_QUERY, &tokenHandle)) {
1259         PyErr_SetFromOSErrnoWithSyscall("OpenProcessToken");
1260         goto error;
1261     }
1262 
1263     CloseHandle(processHandle);
1264     processHandle = NULL;
1265 
1266     // Get the user SID.
1267     bufferSize = 0x100;
1268     while (1) {
1269         user = malloc(bufferSize);
1270         if (user == NULL) {
1271             PyErr_NoMemory();
1272             goto error;
1273         }
1274         if (!GetTokenInformation(tokenHandle, TokenUser, user, bufferSize,
1275                                  &bufferSize))
1276         {
1277             if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
1278                 free(user);
1279                 continue;
1280             }
1281             else {
1282                 PyErr_SetFromOSErrnoWithSyscall("GetTokenInformation");
1283                 goto error;
1284             }
1285         }
1286         break;
1287     }
1288 
1289     CloseHandle(tokenHandle);
1290     tokenHandle = NULL;
1291 
1292     // resolve the SID to a name
1293     nameSize = 0x100;
1294     domainNameSize = 0x100;
1295     while (1) {
1296         name = malloc(nameSize * sizeof(WCHAR));
1297         if (name == NULL) {
1298             PyErr_NoMemory();
1299             goto error;
1300         }
1301         domainName = malloc(domainNameSize * sizeof(WCHAR));
1302         if (domainName == NULL) {
1303             PyErr_NoMemory();
1304             goto error;
1305         }
1306         if (!LookupAccountSidW(NULL, user->User.Sid, name, &nameSize,
1307                                domainName, &domainNameSize, &nameUse))
1308         {
1309             if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
1310                 free(name);
1311                 free(domainName);
1312                 continue;
1313             }
1314             else {
1315                 PyErr_SetFromOSErrnoWithSyscall("LookupAccountSidW");
1316                 goto error;
1317             }
1318         }
1319         break;
1320     }
1321 
1322     py_domain = PyUnicode_FromWideChar(domainName, wcslen(domainName));
1323     if (! py_domain)
1324         goto error;
1325     py_username = PyUnicode_FromWideChar(name, wcslen(name));
1326     if (! py_username)
1327         goto error;
1328     py_tuple = Py_BuildValue("OO", py_domain, py_username);
1329     if (! py_tuple)
1330         goto error;
1331     Py_DECREF(py_domain);
1332     Py_DECREF(py_username);
1333 
1334     free(name);
1335     free(domainName);
1336     free(user);
1337 
1338     return py_tuple;
1339 
1340 error:
1341     if (processHandle != NULL)
1342         CloseHandle(processHandle);
1343     if (tokenHandle != NULL)
1344         CloseHandle(tokenHandle);
1345     if (name != NULL)
1346         free(name);
1347     if (domainName != NULL)
1348         free(domainName);
1349     if (user != NULL)
1350         free(user);
1351     Py_XDECREF(py_domain);
1352     Py_XDECREF(py_username);
1353     Py_XDECREF(py_tuple);
1354     return NULL;
1355 }
1356 
1357 
1358 // https://msdn.microsoft.com/library/aa365928.aspx
1359 // TODO properly handle return code
__GetExtendedTcpTable(_GetExtendedTcpTable call,ULONG address_family,PVOID * data,DWORD * size)1360 static DWORD __GetExtendedTcpTable(_GetExtendedTcpTable call,
1361                                    ULONG address_family,
1362                                    PVOID * data, DWORD * size)
1363 {
1364     // Due to other processes being active on the machine, it's possible
1365     // that the size of the table increases between the moment where we
1366     // query the size and the moment where we query the data.  Therefore, it's
1367     // important to call this in a loop to retry if that happens.
1368     // See https://github.com/giampaolo/psutil/pull/1335 concerning 0xC0000001 error
1369     // and https://github.com/giampaolo/psutil/issues/1294
1370     DWORD error = ERROR_INSUFFICIENT_BUFFER;
1371     *size = 0;
1372     *data = NULL;
1373     error = call(NULL, size, FALSE, address_family,
1374                  TCP_TABLE_OWNER_PID_ALL, 0);
1375     while (error == ERROR_INSUFFICIENT_BUFFER || error == 0xC0000001)
1376     {
1377         *data = malloc(*size);
1378         if (*data == NULL) {
1379             error = ERROR_NOT_ENOUGH_MEMORY;
1380             continue;
1381         }
1382         error = call(*data, size, FALSE, address_family,
1383                      TCP_TABLE_OWNER_PID_ALL, 0);
1384         if (error != NO_ERROR) {
1385             free(*data);
1386             *data = NULL;
1387         }
1388     }
1389     return error;
1390 }
1391 
1392 
1393 // https://msdn.microsoft.com/library/aa365930.aspx
1394 // TODO properly check return value
__GetExtendedUdpTable(_GetExtendedUdpTable call,ULONG address_family,PVOID * data,DWORD * size)1395 static DWORD __GetExtendedUdpTable(_GetExtendedUdpTable call,
1396                                    ULONG address_family,
1397                                    PVOID * data, DWORD * size)
1398 {
1399     // Due to other processes being active on the machine, it's possible
1400     // that the size of the table increases between the moment where we
1401     // query the size and the moment where we query the data.  Therefore, it's
1402     // important to call this in a loop to retry if that happens.
1403     // See https://github.com/giampaolo/psutil/pull/1335 concerning 0xC0000001 error
1404     // and https://github.com/giampaolo/psutil/issues/1294
1405     DWORD error = ERROR_INSUFFICIENT_BUFFER;
1406     *size = 0;
1407     *data = NULL;
1408     error = call(NULL, size, FALSE, address_family,
1409                  UDP_TABLE_OWNER_PID, 0);
1410     while (error == ERROR_INSUFFICIENT_BUFFER || error == 0xC0000001)
1411     {
1412         *data = malloc(*size);
1413         if (*data == NULL) {
1414             error = ERROR_NOT_ENOUGH_MEMORY;
1415             continue;
1416         }
1417         error = call(*data, size, FALSE, address_family,
1418                      UDP_TABLE_OWNER_PID, 0);
1419         if (error != NO_ERROR) {
1420             free(*data);
1421             *data = NULL;
1422         }
1423     }
1424 
1425     if (error == ERROR_NOT_ENOUGH_MEMORY) {
1426         PyErr_NoMemory();
1427         return 1;
1428     }
1429     if (error != NO_ERROR) {
1430         PyErr_SetFromWindowsErr(error);
1431         return 1;
1432     }
1433     return 0;
1434 }
1435 
1436 
1437 #define psutil_conn_decref_objs() \
1438     Py_DECREF(_AF_INET); \
1439     Py_DECREF(_AF_INET6);\
1440     Py_DECREF(_SOCK_STREAM);\
1441     Py_DECREF(_SOCK_DGRAM);
1442 
1443 
1444 /*
1445  * Return a list of network connections opened by a process
1446  */
1447 static PyObject *
psutil_net_connections(PyObject * self,PyObject * args)1448 psutil_net_connections(PyObject *self, PyObject *args) {
1449     static long null_address[4] = { 0, 0, 0, 0 };
1450     unsigned long pid;
1451     int pid_return;
1452     PVOID table = NULL;
1453     DWORD tableSize;
1454     DWORD error;
1455     PMIB_TCPTABLE_OWNER_PID tcp4Table;
1456     PMIB_UDPTABLE_OWNER_PID udp4Table;
1457     PMIB_TCP6TABLE_OWNER_PID tcp6Table;
1458     PMIB_UDP6TABLE_OWNER_PID udp6Table;
1459     ULONG i;
1460     CHAR addressBufferLocal[65];
1461     CHAR addressBufferRemote[65];
1462 
1463     PyObject *py_retlist;
1464     PyObject *py_conn_tuple = NULL;
1465     PyObject *py_af_filter = NULL;
1466     PyObject *py_type_filter = NULL;
1467     PyObject *py_addr_tuple_local = NULL;
1468     PyObject *py_addr_tuple_remote = NULL;
1469     PyObject *_AF_INET = PyLong_FromLong((long)AF_INET);
1470     PyObject *_AF_INET6 = PyLong_FromLong((long)AF_INET6);
1471     PyObject *_SOCK_STREAM = PyLong_FromLong((long)SOCK_STREAM);
1472     PyObject *_SOCK_DGRAM = PyLong_FromLong((long)SOCK_DGRAM);
1473 
1474     // Import some functions.
1475     if (! PyArg_ParseTuple(args, "lOO", &pid, &py_af_filter, &py_type_filter))
1476         goto error;
1477 
1478     if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) {
1479         psutil_conn_decref_objs();
1480         PyErr_SetString(PyExc_TypeError, "arg 2 or 3 is not a sequence");
1481         return NULL;
1482     }
1483 
1484     if (pid != -1) {
1485         pid_return = psutil_pid_is_running(pid);
1486         if (pid_return == 0) {
1487             psutil_conn_decref_objs();
1488             return NoSuchProcess("");
1489         }
1490         else if (pid_return == -1) {
1491             psutil_conn_decref_objs();
1492             return NULL;
1493         }
1494     }
1495 
1496     py_retlist = PyList_New(0);
1497     if (py_retlist == NULL) {
1498         psutil_conn_decref_objs();
1499         return NULL;
1500     }
1501 
1502     // TCP IPv4
1503 
1504     if ((PySequence_Contains(py_af_filter, _AF_INET) == 1) &&
1505             (PySequence_Contains(py_type_filter, _SOCK_STREAM) == 1))
1506     {
1507         table = NULL;
1508         py_conn_tuple = NULL;
1509         py_addr_tuple_local = NULL;
1510         py_addr_tuple_remote = NULL;
1511         tableSize = 0;
1512 
1513         error = __GetExtendedTcpTable(psutil_GetExtendedTcpTable,
1514                                       AF_INET, &table, &tableSize);
1515         if (error != 0)
1516             goto error;
1517         tcp4Table = table;
1518         for (i = 0; i < tcp4Table->dwNumEntries; i++) {
1519             if (pid != -1) {
1520                 if (tcp4Table->table[i].dwOwningPid != pid) {
1521                     continue;
1522                 }
1523             }
1524 
1525             if (tcp4Table->table[i].dwLocalAddr != 0 ||
1526                     tcp4Table->table[i].dwLocalPort != 0)
1527             {
1528                 struct in_addr addr;
1529 
1530                 addr.S_un.S_addr = tcp4Table->table[i].dwLocalAddr;
1531                 psutil_rtlIpv4AddressToStringA(&addr, addressBufferLocal);
1532                 py_addr_tuple_local = Py_BuildValue(
1533                     "(si)",
1534                     addressBufferLocal,
1535                     BYTESWAP_USHORT(tcp4Table->table[i].dwLocalPort));
1536             }
1537             else {
1538                 py_addr_tuple_local = PyTuple_New(0);
1539             }
1540 
1541             if (py_addr_tuple_local == NULL)
1542                 goto error;
1543 
1544             // On Windows <= XP, remote addr is filled even if socket
1545             // is in LISTEN mode in which case we just ignore it.
1546             if ((tcp4Table->table[i].dwRemoteAddr != 0 ||
1547                     tcp4Table->table[i].dwRemotePort != 0) &&
1548                     (tcp4Table->table[i].dwState != MIB_TCP_STATE_LISTEN))
1549             {
1550                 struct in_addr addr;
1551 
1552                 addr.S_un.S_addr = tcp4Table->table[i].dwRemoteAddr;
1553                 psutil_rtlIpv4AddressToStringA(&addr, addressBufferRemote);
1554                 py_addr_tuple_remote = Py_BuildValue(
1555                     "(si)",
1556                     addressBufferRemote,
1557                     BYTESWAP_USHORT(tcp4Table->table[i].dwRemotePort));
1558             }
1559             else
1560             {
1561                 py_addr_tuple_remote = PyTuple_New(0);
1562             }
1563 
1564             if (py_addr_tuple_remote == NULL)
1565                 goto error;
1566 
1567             py_conn_tuple = Py_BuildValue(
1568                 "(iiiNNiI)",
1569                 -1,
1570                 AF_INET,
1571                 SOCK_STREAM,
1572                 py_addr_tuple_local,
1573                 py_addr_tuple_remote,
1574                 tcp4Table->table[i].dwState,
1575                 tcp4Table->table[i].dwOwningPid);
1576             if (!py_conn_tuple)
1577                 goto error;
1578             if (PyList_Append(py_retlist, py_conn_tuple))
1579                 goto error;
1580             Py_CLEAR(py_conn_tuple);
1581         }
1582 
1583         free(table);
1584         table = NULL;
1585         tableSize = 0;
1586     }
1587 
1588     // TCP IPv6
1589     if ((PySequence_Contains(py_af_filter, _AF_INET6) == 1) &&
1590             (PySequence_Contains(py_type_filter, _SOCK_STREAM) == 1) &&
1591             (psutil_rtlIpv6AddressToStringA != NULL))
1592     {
1593         table = NULL;
1594         py_conn_tuple = NULL;
1595         py_addr_tuple_local = NULL;
1596         py_addr_tuple_remote = NULL;
1597         tableSize = 0;
1598 
1599         error = __GetExtendedTcpTable(psutil_GetExtendedTcpTable,
1600                                       AF_INET6, &table, &tableSize);
1601         if (error != 0)
1602             goto error;
1603         tcp6Table = table;
1604         for (i = 0; i < tcp6Table->dwNumEntries; i++)
1605         {
1606             if (pid != -1) {
1607                 if (tcp6Table->table[i].dwOwningPid != pid) {
1608                     continue;
1609                 }
1610             }
1611 
1612             if (memcmp(tcp6Table->table[i].ucLocalAddr, null_address, 16)
1613                     != 0 || tcp6Table->table[i].dwLocalPort != 0)
1614             {
1615                 struct in6_addr addr;
1616 
1617                 memcpy(&addr, tcp6Table->table[i].ucLocalAddr, 16);
1618                 psutil_rtlIpv6AddressToStringA(&addr, addressBufferLocal);
1619                 py_addr_tuple_local = Py_BuildValue(
1620                     "(si)",
1621                     addressBufferLocal,
1622                     BYTESWAP_USHORT(tcp6Table->table[i].dwLocalPort));
1623             }
1624             else {
1625                 py_addr_tuple_local = PyTuple_New(0);
1626             }
1627 
1628             if (py_addr_tuple_local == NULL)
1629                 goto error;
1630 
1631             // On Windows <= XP, remote addr is filled even if socket
1632             // is in LISTEN mode in which case we just ignore it.
1633             if ((memcmp(tcp6Table->table[i].ucRemoteAddr, null_address, 16)
1634                     != 0 ||
1635                     tcp6Table->table[i].dwRemotePort != 0) &&
1636                     (tcp6Table->table[i].dwState != MIB_TCP_STATE_LISTEN))
1637             {
1638                 struct in6_addr addr;
1639 
1640                 memcpy(&addr, tcp6Table->table[i].ucRemoteAddr, 16);
1641                 psutil_rtlIpv6AddressToStringA(&addr, addressBufferRemote);
1642                 py_addr_tuple_remote = Py_BuildValue(
1643                     "(si)",
1644                     addressBufferRemote,
1645                     BYTESWAP_USHORT(tcp6Table->table[i].dwRemotePort));
1646             }
1647             else {
1648                 py_addr_tuple_remote = PyTuple_New(0);
1649             }
1650 
1651             if (py_addr_tuple_remote == NULL)
1652                 goto error;
1653 
1654             py_conn_tuple = Py_BuildValue(
1655                 "(iiiNNiI)",
1656                 -1,
1657                 AF_INET6,
1658                 SOCK_STREAM,
1659                 py_addr_tuple_local,
1660                 py_addr_tuple_remote,
1661                 tcp6Table->table[i].dwState,
1662                 tcp6Table->table[i].dwOwningPid);
1663             if (!py_conn_tuple)
1664                 goto error;
1665             if (PyList_Append(py_retlist, py_conn_tuple))
1666                 goto error;
1667             Py_CLEAR(py_conn_tuple);
1668         }
1669 
1670         free(table);
1671         table = NULL;
1672         tableSize = 0;
1673     }
1674 
1675     // UDP IPv4
1676 
1677     if ((PySequence_Contains(py_af_filter, _AF_INET) == 1) &&
1678             (PySequence_Contains(py_type_filter, _SOCK_DGRAM) == 1))
1679     {
1680         table = NULL;
1681         py_conn_tuple = NULL;
1682         py_addr_tuple_local = NULL;
1683         py_addr_tuple_remote = NULL;
1684         tableSize = 0;
1685         error = __GetExtendedUdpTable(psutil_GetExtendedUdpTable,
1686                                       AF_INET, &table, &tableSize);
1687         if (error != 0)
1688             goto error;
1689         udp4Table = table;
1690         for (i = 0; i < udp4Table->dwNumEntries; i++)
1691         {
1692             if (pid != -1) {
1693                 if (udp4Table->table[i].dwOwningPid != pid) {
1694                     continue;
1695                 }
1696             }
1697 
1698             if (udp4Table->table[i].dwLocalAddr != 0 ||
1699                 udp4Table->table[i].dwLocalPort != 0)
1700             {
1701                 struct in_addr addr;
1702 
1703                 addr.S_un.S_addr = udp4Table->table[i].dwLocalAddr;
1704                 psutil_rtlIpv4AddressToStringA(&addr, addressBufferLocal);
1705                 py_addr_tuple_local = Py_BuildValue(
1706                     "(si)",
1707                     addressBufferLocal,
1708                     BYTESWAP_USHORT(udp4Table->table[i].dwLocalPort));
1709             }
1710             else {
1711                 py_addr_tuple_local = PyTuple_New(0);
1712             }
1713 
1714             if (py_addr_tuple_local == NULL)
1715                 goto error;
1716 
1717             py_conn_tuple = Py_BuildValue(
1718                 "(iiiNNiI)",
1719                 -1,
1720                 AF_INET,
1721                 SOCK_DGRAM,
1722                 py_addr_tuple_local,
1723                 PyTuple_New(0),
1724                 PSUTIL_CONN_NONE,
1725                 udp4Table->table[i].dwOwningPid);
1726             if (!py_conn_tuple)
1727                 goto error;
1728             if (PyList_Append(py_retlist, py_conn_tuple))
1729                 goto error;
1730             Py_CLEAR(py_conn_tuple);
1731         }
1732 
1733         free(table);
1734         table = NULL;
1735         tableSize = 0;
1736     }
1737 
1738     // UDP IPv6
1739 
1740     if ((PySequence_Contains(py_af_filter, _AF_INET6) == 1) &&
1741             (PySequence_Contains(py_type_filter, _SOCK_DGRAM) == 1) &&
1742             (psutil_rtlIpv6AddressToStringA != NULL))
1743     {
1744         table = NULL;
1745         py_conn_tuple = NULL;
1746         py_addr_tuple_local = NULL;
1747         py_addr_tuple_remote = NULL;
1748         tableSize = 0;
1749         error = __GetExtendedUdpTable(psutil_GetExtendedUdpTable,
1750                                       AF_INET6, &table, &tableSize);
1751         if (error != 0)
1752             goto error;
1753         udp6Table = table;
1754         for (i = 0; i < udp6Table->dwNumEntries; i++) {
1755             if (pid != -1) {
1756                 if (udp6Table->table[i].dwOwningPid != pid) {
1757                     continue;
1758                 }
1759             }
1760 
1761             if (memcmp(udp6Table->table[i].ucLocalAddr, null_address, 16)
1762                     != 0 || udp6Table->table[i].dwLocalPort != 0)
1763             {
1764                 struct in6_addr addr;
1765 
1766                 memcpy(&addr, udp6Table->table[i].ucLocalAddr, 16);
1767                 psutil_rtlIpv6AddressToStringA(&addr, addressBufferLocal);
1768                 py_addr_tuple_local = Py_BuildValue(
1769                     "(si)",
1770                     addressBufferLocal,
1771                     BYTESWAP_USHORT(udp6Table->table[i].dwLocalPort));
1772             }
1773             else {
1774                 py_addr_tuple_local = PyTuple_New(0);
1775             }
1776 
1777             if (py_addr_tuple_local == NULL)
1778                 goto error;
1779 
1780             py_conn_tuple = Py_BuildValue(
1781                 "(iiiNNiI)",
1782                 -1,
1783                 AF_INET6,
1784                 SOCK_DGRAM,
1785                 py_addr_tuple_local,
1786                 PyTuple_New(0),
1787                 PSUTIL_CONN_NONE,
1788                 udp6Table->table[i].dwOwningPid);
1789             if (!py_conn_tuple)
1790                 goto error;
1791             if (PyList_Append(py_retlist, py_conn_tuple))
1792                 goto error;
1793             Py_CLEAR(py_conn_tuple);
1794         }
1795 
1796         free(table);
1797         table = NULL;
1798         tableSize = 0;
1799     }
1800 
1801     psutil_conn_decref_objs();
1802     return py_retlist;
1803 
1804 error:
1805     psutil_conn_decref_objs();
1806     Py_XDECREF(py_conn_tuple);
1807     Py_XDECREF(py_addr_tuple_local);
1808     Py_XDECREF(py_addr_tuple_remote);
1809     Py_DECREF(py_retlist);
1810     if (table != NULL)
1811         free(table);
1812     return NULL;
1813 }
1814 
1815 
1816 /*
1817  * Get process priority as a Python integer.
1818  */
1819 static PyObject *
psutil_proc_priority_get(PyObject * self,PyObject * args)1820 psutil_proc_priority_get(PyObject *self, PyObject *args) {
1821     long pid;
1822     DWORD priority;
1823     HANDLE hProcess;
1824 
1825     if (! PyArg_ParseTuple(args, "l", &pid))
1826         return NULL;
1827 
1828     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
1829     if (hProcess == NULL)
1830         return NULL;
1831 
1832     priority = GetPriorityClass(hProcess);
1833     if (priority == 0) {
1834         PyErr_SetFromWindowsErr(0);
1835         CloseHandle(hProcess);
1836         return NULL;
1837     }
1838     CloseHandle(hProcess);
1839     return Py_BuildValue("i", priority);
1840 }
1841 
1842 
1843 /*
1844  * Set process priority.
1845  */
1846 static PyObject *
psutil_proc_priority_set(PyObject * self,PyObject * args)1847 psutil_proc_priority_set(PyObject *self, PyObject *args) {
1848     long pid;
1849     int priority;
1850     int retval;
1851     HANDLE hProcess;
1852     DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION;
1853 
1854     if (! PyArg_ParseTuple(args, "li", &pid, &priority))
1855         return NULL;
1856     hProcess = psutil_handle_from_pid(pid, access);
1857     if (hProcess == NULL)
1858         return NULL;
1859 
1860     retval = SetPriorityClass(hProcess, priority);
1861     if (retval == 0) {
1862         PyErr_SetFromWindowsErr(0);
1863         CloseHandle(hProcess);
1864         return NULL;
1865     }
1866 
1867     CloseHandle(hProcess);
1868     Py_RETURN_NONE;
1869 }
1870 
1871 
1872 #if (_WIN32_WINNT >= 0x0600)  // Windows Vista
1873 /*
1874  * Get process IO priority as a Python integer.
1875  */
1876 static PyObject *
psutil_proc_io_priority_get(PyObject * self,PyObject * args)1877 psutil_proc_io_priority_get(PyObject *self, PyObject *args) {
1878     long pid;
1879     HANDLE hProcess;
1880     DWORD IoPriority;
1881     NTSTATUS status;
1882 
1883     if (! PyArg_ParseTuple(args, "l", &pid))
1884         return NULL;
1885 
1886     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
1887     if (hProcess == NULL)
1888         return NULL;
1889 
1890     status = psutil_NtQueryInformationProcess(
1891         hProcess,
1892         ProcessIoPriority,
1893         &IoPriority,
1894         sizeof(DWORD),
1895         NULL
1896     );
1897 
1898     CloseHandle(hProcess);
1899     if (! NT_SUCCESS(status))
1900         return psutil_SetFromNTStatusErr(status, "NtQueryInformationProcess");
1901     return Py_BuildValue("i", IoPriority);
1902 }
1903 
1904 
1905 /*
1906  * Set process IO priority.
1907  */
1908 static PyObject *
psutil_proc_io_priority_set(PyObject * self,PyObject * args)1909 psutil_proc_io_priority_set(PyObject *self, PyObject *args) {
1910     long pid;
1911     DWORD prio;
1912     HANDLE hProcess;
1913     NTSTATUS status;
1914     DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION;
1915 
1916     if (! PyArg_ParseTuple(args, "li", &pid, &prio))
1917         return NULL;
1918 
1919     hProcess = psutil_handle_from_pid(pid, access);
1920     if (hProcess == NULL)
1921         return NULL;
1922 
1923     status = psutil_NtSetInformationProcess(
1924         hProcess,
1925         ProcessIoPriority,
1926         (PVOID)&prio,
1927         sizeof(DWORD)
1928     );
1929 
1930     CloseHandle(hProcess);
1931     if (! NT_SUCCESS(status))
1932         return psutil_SetFromNTStatusErr(status, "NtSetInformationProcess");
1933     Py_RETURN_NONE;
1934 }
1935 #endif
1936 
1937 
1938 /*
1939  * Return a Python tuple referencing process I/O counters.
1940  */
1941 static PyObject *
psutil_proc_io_counters(PyObject * self,PyObject * args)1942 psutil_proc_io_counters(PyObject *self, PyObject *args) {
1943     DWORD pid;
1944     HANDLE hProcess;
1945     IO_COUNTERS IoCounters;
1946 
1947     if (! PyArg_ParseTuple(args, "l", &pid))
1948         return NULL;
1949     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
1950     if (NULL == hProcess)
1951         return NULL;
1952 
1953     if (! GetProcessIoCounters(hProcess, &IoCounters)) {
1954         PyErr_SetFromWindowsErr(0);
1955         CloseHandle(hProcess);
1956         return NULL;
1957     }
1958 
1959     CloseHandle(hProcess);
1960     return Py_BuildValue("(KKKKKK)",
1961                          IoCounters.ReadOperationCount,
1962                          IoCounters.WriteOperationCount,
1963                          IoCounters.ReadTransferCount,
1964                          IoCounters.WriteTransferCount,
1965                          IoCounters.OtherOperationCount,
1966                          IoCounters.OtherTransferCount);
1967 }
1968 
1969 
1970 /*
1971  * Return process CPU affinity as a bitmask
1972  */
1973 static PyObject *
psutil_proc_cpu_affinity_get(PyObject * self,PyObject * args)1974 psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
1975     DWORD pid;
1976     HANDLE hProcess;
1977     DWORD_PTR proc_mask;
1978     DWORD_PTR system_mask;
1979 
1980     if (! PyArg_ParseTuple(args, "l", &pid))
1981         return NULL;
1982     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
1983     if (hProcess == NULL) {
1984         return NULL;
1985     }
1986     if (GetProcessAffinityMask(hProcess, &proc_mask, &system_mask) == 0) {
1987         PyErr_SetFromWindowsErr(0);
1988         CloseHandle(hProcess);
1989         return NULL;
1990     }
1991 
1992     CloseHandle(hProcess);
1993 #ifdef _WIN64
1994     return Py_BuildValue("K", (unsigned long long)proc_mask);
1995 #else
1996     return Py_BuildValue("k", (unsigned long)proc_mask);
1997 #endif
1998 }
1999 
2000 
2001 /*
2002  * Set process CPU affinity
2003  */
2004 static PyObject *
psutil_proc_cpu_affinity_set(PyObject * self,PyObject * args)2005 psutil_proc_cpu_affinity_set(PyObject *self, PyObject *args) {
2006     DWORD pid;
2007     HANDLE hProcess;
2008     DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION;
2009     DWORD_PTR mask;
2010 
2011 #ifdef _WIN64
2012     if (! PyArg_ParseTuple(args, "lK", &pid, &mask))
2013 #else
2014     if (! PyArg_ParseTuple(args, "lk", &pid, &mask))
2015 #endif
2016     {
2017         return NULL;
2018     }
2019     hProcess = psutil_handle_from_pid(pid, access);
2020     if (hProcess == NULL)
2021         return NULL;
2022 
2023     if (SetProcessAffinityMask(hProcess, mask) == 0) {
2024         PyErr_SetFromWindowsErr(0);
2025         CloseHandle(hProcess);
2026         return NULL;
2027     }
2028 
2029     CloseHandle(hProcess);
2030     Py_RETURN_NONE;
2031 }
2032 
2033 
2034 /*
2035  * Return True if all process threads are in waiting/suspended state.
2036  */
2037 static PyObject *
psutil_proc_is_suspended(PyObject * self,PyObject * args)2038 psutil_proc_is_suspended(PyObject *self, PyObject *args) {
2039     DWORD pid;
2040     ULONG i;
2041     PSYSTEM_PROCESS_INFORMATION process;
2042     PVOID buffer;
2043 
2044     if (! PyArg_ParseTuple(args, "l", &pid))
2045         return NULL;
2046     if (! psutil_get_proc_info(pid, &process, &buffer))
2047         return NULL;
2048     for (i = 0; i < process->NumberOfThreads; i++) {
2049         if (process->Threads[i].ThreadState != Waiting ||
2050                 process->Threads[i].WaitReason != Suspended)
2051         {
2052             free(buffer);
2053             Py_RETURN_FALSE;
2054         }
2055     }
2056     free(buffer);
2057     Py_RETURN_TRUE;
2058 }
2059 
2060 
2061 /*
2062  * Return path's disk total and free as a Python tuple.
2063  */
2064 static PyObject *
psutil_disk_usage(PyObject * self,PyObject * args)2065 psutil_disk_usage(PyObject *self, PyObject *args) {
2066     BOOL retval;
2067     ULARGE_INTEGER _, total, free;
2068     char *path;
2069 
2070     if (PyArg_ParseTuple(args, "u", &path)) {
2071         Py_BEGIN_ALLOW_THREADS
2072         retval = GetDiskFreeSpaceExW((LPCWSTR)path, &_, &total, &free);
2073         Py_END_ALLOW_THREADS
2074         goto return_;
2075     }
2076 
2077     // on Python 2 we also want to accept plain strings other
2078     // than Unicode
2079 #if PY_MAJOR_VERSION <= 2
2080     PyErr_Clear();  // drop the argument parsing error
2081     if (PyArg_ParseTuple(args, "s", &path)) {
2082         Py_BEGIN_ALLOW_THREADS
2083         retval = GetDiskFreeSpaceEx(path, &_, &total, &free);
2084         Py_END_ALLOW_THREADS
2085         goto return_;
2086     }
2087 #endif
2088 
2089     return NULL;
2090 
2091 return_:
2092     if (retval == 0)
2093         return PyErr_SetFromWindowsErrWithFilename(0, path);
2094     else
2095         return Py_BuildValue("(LL)", total.QuadPart, free.QuadPart);
2096 }
2097 
2098 
2099 /*
2100  * Return a Python list of named tuples with overall network I/O information
2101  */
2102 static PyObject *
psutil_net_io_counters(PyObject * self,PyObject * args)2103 psutil_net_io_counters(PyObject *self, PyObject *args) {
2104     DWORD dwRetVal = 0;
2105 
2106 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
2107     MIB_IF_ROW2 *pIfRow = NULL;
2108 #else // Windows XP
2109     MIB_IFROW *pIfRow = NULL;
2110 #endif
2111 
2112     PIP_ADAPTER_ADDRESSES pAddresses = NULL;
2113     PIP_ADAPTER_ADDRESSES pCurrAddresses = NULL;
2114     PyObject *py_retdict = PyDict_New();
2115     PyObject *py_nic_info = NULL;
2116     PyObject *py_nic_name = NULL;
2117 
2118     if (py_retdict == NULL)
2119         return NULL;
2120     pAddresses = psutil_get_nic_addresses();
2121     if (pAddresses == NULL)
2122         goto error;
2123     pCurrAddresses = pAddresses;
2124 
2125     while (pCurrAddresses) {
2126         py_nic_name = NULL;
2127         py_nic_info = NULL;
2128 
2129 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
2130         pIfRow = (MIB_IF_ROW2 *) malloc(sizeof(MIB_IF_ROW2));
2131 #else // Windows XP
2132         pIfRow = (MIB_IFROW *) malloc(sizeof(MIB_IFROW));
2133 #endif
2134 
2135         if (pIfRow == NULL) {
2136             PyErr_NoMemory();
2137             goto error;
2138         }
2139 
2140 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
2141         SecureZeroMemory((PVOID)pIfRow, sizeof(MIB_IF_ROW2));
2142         pIfRow->InterfaceIndex = pCurrAddresses->IfIndex;
2143         dwRetVal = GetIfEntry2(pIfRow);
2144 #else // Windows XP
2145         pIfRow->dwIndex = pCurrAddresses->IfIndex;
2146         dwRetVal = GetIfEntry(pIfRow);
2147 #endif
2148 
2149         if (dwRetVal != NO_ERROR) {
2150             PyErr_SetString(PyExc_RuntimeError,
2151                             "GetIfEntry() or GetIfEntry2() syscalls failed.");
2152             goto error;
2153         }
2154 
2155 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
2156         py_nic_info = Py_BuildValue("(KKKKKKKK)",
2157                                     pIfRow->OutOctets,
2158                                     pIfRow->InOctets,
2159                                     (pIfRow->OutUcastPkts + pIfRow->OutNUcastPkts),
2160                                     (pIfRow->InUcastPkts + pIfRow->InNUcastPkts),
2161                                     pIfRow->InErrors,
2162                                     pIfRow->OutErrors,
2163                                     pIfRow->InDiscards,
2164                                     pIfRow->OutDiscards);
2165 #else // Windows XP
2166         py_nic_info = Py_BuildValue("(kkkkkkkk)",
2167                                     pIfRow->dwOutOctets,
2168                                     pIfRow->dwInOctets,
2169                                     (pIfRow->dwOutUcastPkts + pIfRow->dwOutNUcastPkts),
2170                                     (pIfRow->dwInUcastPkts + pIfRow->dwInNUcastPkts),
2171                                     pIfRow->dwInErrors,
2172                                     pIfRow->dwOutErrors,
2173                                     pIfRow->dwInDiscards,
2174                                     pIfRow->dwOutDiscards);
2175 #endif
2176 
2177         if (!py_nic_info)
2178             goto error;
2179 
2180         py_nic_name = PyUnicode_FromWideChar(
2181             pCurrAddresses->FriendlyName,
2182             wcslen(pCurrAddresses->FriendlyName));
2183 
2184         if (py_nic_name == NULL)
2185             goto error;
2186         if (PyDict_SetItem(py_retdict, py_nic_name, py_nic_info))
2187             goto error;
2188         Py_CLEAR(py_nic_name);
2189         Py_CLEAR(py_nic_info);
2190 
2191         free(pIfRow);
2192         pCurrAddresses = pCurrAddresses->Next;
2193     }
2194 
2195     free(pAddresses);
2196     return py_retdict;
2197 
2198 error:
2199     Py_XDECREF(py_nic_name);
2200     Py_XDECREF(py_nic_info);
2201     Py_DECREF(py_retdict);
2202     if (pAddresses != NULL)
2203         free(pAddresses);
2204     if (pIfRow != NULL)
2205         free(pIfRow);
2206     return NULL;
2207 }
2208 
2209 
2210 /*
2211  * Return a Python dict of tuples for disk I/O information. This may
2212  * require running "diskperf -y" command first.
2213  */
2214 static PyObject *
psutil_disk_io_counters(PyObject * self,PyObject * args)2215 psutil_disk_io_counters(PyObject *self, PyObject *args) {
2216     DISK_PERFORMANCE diskPerformance;
2217     DWORD dwSize;
2218     HANDLE hDevice = NULL;
2219     char szDevice[MAX_PATH];
2220     char szDeviceDisplay[MAX_PATH];
2221     int devNum;
2222     int i;
2223     DWORD ioctrlSize;
2224     BOOL ret;
2225     PyObject *py_retdict = PyDict_New();
2226     PyObject *py_tuple = NULL;
2227 
2228     if (py_retdict == NULL)
2229         return NULL;
2230     // Apparently there's no way to figure out how many times we have
2231     // to iterate in order to find valid drives.
2232     // Let's assume 32, which is higher than 26, the number of letters
2233     // in the alphabet (from A:\ to Z:\).
2234     for (devNum = 0; devNum <= 32; ++devNum) {
2235         py_tuple = NULL;
2236         sprintf_s(szDevice, MAX_PATH, "\\\\.\\PhysicalDrive%d", devNum);
2237         hDevice = CreateFile(szDevice, 0, FILE_SHARE_READ | FILE_SHARE_WRITE,
2238                              NULL, OPEN_EXISTING, 0, NULL);
2239         if (hDevice == INVALID_HANDLE_VALUE)
2240             continue;
2241 
2242         // DeviceIoControl() sucks!
2243         i = 0;
2244         ioctrlSize = sizeof(diskPerformance);
2245         while (1) {
2246             i += 1;
2247             ret = DeviceIoControl(
2248                 hDevice, IOCTL_DISK_PERFORMANCE, NULL, 0, &diskPerformance,
2249                 ioctrlSize, &dwSize, NULL);
2250             if (ret != 0)
2251                 break;  // OK!
2252             if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
2253                 // Retry with a bigger buffer (+ limit for retries).
2254                 if (i <= 1024) {
2255                     ioctrlSize *= 2;
2256                     continue;
2257                 }
2258             }
2259             else if (GetLastError() == ERROR_INVALID_FUNCTION) {
2260                 // This happens on AppVeyor:
2261                 // https://ci.appveyor.com/project/giampaolo/psutil/build/
2262                 //      1364/job/ascpdi271b06jle3
2263                 // Assume it means we're dealing with some exotic disk
2264                 // and go on.
2265                 psutil_debug("DeviceIoControl -> ERROR_INVALID_FUNCTION; "
2266                              "ignore PhysicalDrive%i", devNum);
2267                 goto next;
2268             }
2269             else if (GetLastError() == ERROR_NOT_SUPPORTED) {
2270                 // Again, let's assume we're dealing with some exotic disk.
2271                 psutil_debug("DeviceIoControl -> ERROR_NOT_SUPPORTED; "
2272                              "ignore PhysicalDrive%i", devNum);
2273                 goto next;
2274             }
2275             // XXX: it seems we should also catch ERROR_INVALID_PARAMETER:
2276             // https://sites.ualberta.ca/dept/aict/uts/software/openbsd/
2277             //     ports/4.1/i386/openafs/w-openafs-1.4.14-transarc/
2278             //     openafs-1.4.14/src/usd/usd_nt.c
2279 
2280             // XXX: we can also bump into ERROR_MORE_DATA in which case
2281             // (quoting doc) we're supposed to retry with a bigger buffer
2282             // and specify  a new "starting point", whatever it means.
2283             PyErr_SetFromWindowsErr(0);
2284             goto error;
2285         }
2286 
2287         sprintf_s(szDeviceDisplay, MAX_PATH, "PhysicalDrive%i", devNum);
2288         py_tuple = Py_BuildValue(
2289             "(IILLKK)",
2290             diskPerformance.ReadCount,
2291             diskPerformance.WriteCount,
2292             diskPerformance.BytesRead,
2293             diskPerformance.BytesWritten,
2294             // convert to ms:
2295             // https://github.com/giampaolo/psutil/issues/1012
2296             (unsigned long long)
2297                 (diskPerformance.ReadTime.QuadPart) / 10000000,
2298             (unsigned long long)
2299                 (diskPerformance.WriteTime.QuadPart) / 10000000);
2300         if (!py_tuple)
2301             goto error;
2302         if (PyDict_SetItemString(py_retdict, szDeviceDisplay, py_tuple))
2303             goto error;
2304         Py_CLEAR(py_tuple);
2305 
2306 next:
2307         CloseHandle(hDevice);
2308     }
2309 
2310     return py_retdict;
2311 
2312 error:
2313     Py_XDECREF(py_tuple);
2314     Py_DECREF(py_retdict);
2315     if (hDevice != NULL)
2316         CloseHandle(hDevice);
2317     return NULL;
2318 }
2319 
2320 
psutil_get_drive_type(int type)2321 static char *psutil_get_drive_type(int type) {
2322     switch (type) {
2323         case DRIVE_FIXED:
2324             return "fixed";
2325         case DRIVE_CDROM:
2326             return "cdrom";
2327         case DRIVE_REMOVABLE:
2328             return "removable";
2329         case DRIVE_UNKNOWN:
2330             return "unknown";
2331         case DRIVE_NO_ROOT_DIR:
2332             return "unmounted";
2333         case DRIVE_REMOTE:
2334             return "remote";
2335         case DRIVE_RAMDISK:
2336             return "ramdisk";
2337         default:
2338             return "?";
2339     }
2340 }
2341 
2342 
2343 #ifndef _ARRAYSIZE
2344 #define _ARRAYSIZE(a) (sizeof(a)/sizeof(a[0]))
2345 #endif
2346 
2347 
2348 /*
2349  * Return disk partitions as a list of tuples such as
2350  * (drive_letter, drive_letter, type, "")
2351  */
2352 static PyObject *
psutil_disk_partitions(PyObject * self,PyObject * args)2353 psutil_disk_partitions(PyObject *self, PyObject *args) {
2354     DWORD num_bytes;
2355     char drive_strings[255];
2356     char *drive_letter = drive_strings;
2357     char mp_buf[MAX_PATH];
2358     char mp_path[MAX_PATH];
2359     int all;
2360     int type;
2361     int ret;
2362     unsigned int old_mode = 0;
2363     char opts[20];
2364     HANDLE mp_h;
2365     BOOL mp_flag= TRUE;
2366     LPTSTR fs_type[MAX_PATH + 1] = { 0 };
2367     DWORD pflags = 0;
2368     PyObject *py_all;
2369     PyObject *py_retlist = PyList_New(0);
2370     PyObject *py_tuple = NULL;
2371 
2372     if (py_retlist == NULL) {
2373         return NULL;
2374     }
2375 
2376     // avoid to visualize a message box in case something goes wrong
2377     // see https://github.com/giampaolo/psutil/issues/264
2378     old_mode = SetErrorMode(SEM_FAILCRITICALERRORS);
2379 
2380     if (! PyArg_ParseTuple(args, "O", &py_all))
2381         goto error;
2382     all = PyObject_IsTrue(py_all);
2383 
2384     Py_BEGIN_ALLOW_THREADS
2385     num_bytes = GetLogicalDriveStrings(254, drive_letter);
2386     Py_END_ALLOW_THREADS
2387 
2388     if (num_bytes == 0) {
2389         PyErr_SetFromWindowsErr(0);
2390         goto error;
2391     }
2392 
2393     while (*drive_letter != 0) {
2394         py_tuple = NULL;
2395         opts[0] = 0;
2396         fs_type[0] = 0;
2397 
2398         Py_BEGIN_ALLOW_THREADS
2399         type = GetDriveType(drive_letter);
2400         Py_END_ALLOW_THREADS
2401 
2402         // by default we only show hard drives and cd-roms
2403         if (all == 0) {
2404             if ((type == DRIVE_UNKNOWN) ||
2405                     (type == DRIVE_NO_ROOT_DIR) ||
2406                     (type == DRIVE_REMOTE) ||
2407                     (type == DRIVE_RAMDISK)) {
2408                 goto next;
2409             }
2410             // floppy disk: skip it by default as it introduces a
2411             // considerable slowdown.
2412             if ((type == DRIVE_REMOVABLE) &&
2413                     (strcmp(drive_letter, "A:\\")  == 0)) {
2414                 goto next;
2415             }
2416         }
2417 
2418         ret = GetVolumeInformation(
2419             (LPCTSTR)drive_letter, NULL, _ARRAYSIZE(drive_letter),
2420             NULL, NULL, &pflags, (LPTSTR)fs_type, _ARRAYSIZE(fs_type));
2421         if (ret == 0) {
2422             // We might get here in case of a floppy hard drive, in
2423             // which case the error is (21, "device not ready").
2424             // Let's pretend it didn't happen as we already have
2425             // the drive name and type ('removable').
2426             strcat_s(opts, _countof(opts), "");
2427             SetLastError(0);
2428         }
2429         else {
2430             if (pflags & FILE_READ_ONLY_VOLUME)
2431                 strcat_s(opts, _countof(opts), "ro");
2432             else
2433                 strcat_s(opts, _countof(opts), "rw");
2434             if (pflags & FILE_VOLUME_IS_COMPRESSED)
2435                 strcat_s(opts, _countof(opts), ",compressed");
2436 
2437             // Check for mount points on this volume and add/get info
2438             // (checks first to know if we can even have mount points)
2439             if (pflags & FILE_SUPPORTS_REPARSE_POINTS) {
2440 
2441                 mp_h = FindFirstVolumeMountPoint(drive_letter, mp_buf, MAX_PATH);
2442                 if (mp_h != INVALID_HANDLE_VALUE) {
2443                     while (mp_flag) {
2444 
2445                         // Append full mount path with drive letter
2446                         strcpy_s(mp_path, _countof(mp_path), drive_letter);
2447                         strcat_s(mp_path, _countof(mp_path), mp_buf);
2448 
2449                         py_tuple = Py_BuildValue(
2450                             "(ssss)",
2451                             drive_letter,
2452                             mp_path,
2453                             fs_type, // Typically NTFS
2454                             opts);
2455 
2456                         if (!py_tuple || PyList_Append(py_retlist, py_tuple) == -1) {
2457                             FindVolumeMountPointClose(mp_h);
2458                             goto error;
2459                         }
2460 
2461                         Py_CLEAR(py_tuple);
2462 
2463                         // Continue looking for more mount points
2464                         mp_flag = FindNextVolumeMountPoint(mp_h, mp_buf, MAX_PATH);
2465                     }
2466                     FindVolumeMountPointClose(mp_h);
2467                 }
2468 
2469             }
2470         }
2471 
2472         if (strlen(opts) > 0)
2473             strcat_s(opts, _countof(opts), ",");
2474         strcat_s(opts, _countof(opts), psutil_get_drive_type(type));
2475 
2476         py_tuple = Py_BuildValue(
2477             "(ssss)",
2478             drive_letter,
2479             drive_letter,
2480             fs_type,  // either FAT, FAT32, NTFS, HPFS, CDFS, UDF or NWFS
2481             opts);
2482         if (!py_tuple)
2483             goto error;
2484         if (PyList_Append(py_retlist, py_tuple))
2485             goto error;
2486         Py_CLEAR(py_tuple);
2487         goto next;
2488 
2489 next:
2490         drive_letter = strchr(drive_letter, 0) + 1;
2491     }
2492 
2493     SetErrorMode(old_mode);
2494     return py_retlist;
2495 
2496 error:
2497     SetErrorMode(old_mode);
2498     Py_XDECREF(py_tuple);
2499     Py_DECREF(py_retlist);
2500     return NULL;
2501 }
2502 
2503 /*
2504  * Return a Python dict of tuples for disk I/O information
2505  */
2506 static PyObject *
psutil_users(PyObject * self,PyObject * args)2507 psutil_users(PyObject *self, PyObject *args) {
2508     HANDLE hServer = WTS_CURRENT_SERVER_HANDLE;
2509     WCHAR *buffer_user = NULL;
2510     LPTSTR buffer_addr = NULL;
2511     PWTS_SESSION_INFO sessions = NULL;
2512     DWORD count;
2513     DWORD i;
2514     DWORD sessionId;
2515     DWORD bytes;
2516     PWTS_CLIENT_ADDRESS address;
2517     char address_str[50];
2518     long long unix_time;
2519     WINSTATION_INFO station_info;
2520     ULONG returnLen;
2521     PyObject *py_tuple = NULL;
2522     PyObject *py_address = NULL;
2523     PyObject *py_username = NULL;
2524     PyObject *py_retlist = PyList_New(0);
2525 
2526     if (py_retlist == NULL)
2527         return NULL;
2528 
2529     if (WTSEnumerateSessions(hServer, 0, 1, &sessions, &count) == 0) {
2530         PyErr_SetFromOSErrnoWithSyscall("WTSEnumerateSessions");
2531         goto error;
2532     }
2533 
2534     for (i = 0; i < count; i++) {
2535         py_address = NULL;
2536         py_tuple = NULL;
2537         sessionId = sessions[i].SessionId;
2538         if (buffer_user != NULL)
2539             WTSFreeMemory(buffer_user);
2540         if (buffer_addr != NULL)
2541             WTSFreeMemory(buffer_addr);
2542 
2543         buffer_user = NULL;
2544         buffer_addr = NULL;
2545 
2546         // username
2547         bytes = 0;
2548         if (WTSQuerySessionInformationW(hServer, sessionId, WTSUserName,
2549                                         &buffer_user, &bytes) == 0) {
2550             PyErr_SetFromOSErrnoWithSyscall("WTSQuerySessionInformationW");
2551             goto error;
2552         }
2553         if (bytes <= 2)
2554             continue;
2555 
2556         // address
2557         bytes = 0;
2558         if (WTSQuerySessionInformation(hServer, sessionId, WTSClientAddress,
2559                                        &buffer_addr, &bytes) == 0) {
2560             PyErr_SetFromOSErrnoWithSyscall("WTSQuerySessionInformation");
2561             goto error;
2562         }
2563 
2564         address = (PWTS_CLIENT_ADDRESS)buffer_addr;
2565         if (address->AddressFamily == 0) {  // AF_INET
2566             sprintf_s(address_str,
2567                       _countof(address_str),
2568                       "%u.%u.%u.%u",
2569                       address->Address[0],
2570                       address->Address[1],
2571                       address->Address[2],
2572                       address->Address[3]);
2573             py_address = Py_BuildValue("s", address_str);
2574             if (!py_address)
2575                 goto error;
2576         }
2577         else {
2578             py_address = Py_None;
2579         }
2580 
2581         // login time
2582         if (! psutil_WinStationQueryInformationW(
2583                 hServer,
2584                 sessionId,
2585                 WinStationInformation,
2586                 &station_info,
2587                 sizeof(station_info),
2588                 &returnLen))
2589         {
2590             PyErr_SetFromOSErrnoWithSyscall("WinStationQueryInformationW");
2591             goto error;
2592         }
2593 
2594         unix_time = ((LONGLONG)station_info.ConnectTime.dwHighDateTime) << 32;
2595         unix_time += \
2596             station_info.ConnectTime.dwLowDateTime - 116444736000000000LL;
2597         unix_time /= 10000000;
2598 
2599         py_username = PyUnicode_FromWideChar(buffer_user, wcslen(buffer_user));
2600         if (py_username == NULL)
2601             goto error;
2602         py_tuple = Py_BuildValue("OOd",
2603                                  py_username,
2604                                  py_address,
2605                                  (double)unix_time);
2606         if (!py_tuple)
2607             goto error;
2608         if (PyList_Append(py_retlist, py_tuple))
2609             goto error;
2610         Py_CLEAR(py_username);
2611         Py_CLEAR(py_address);
2612         Py_CLEAR(py_tuple);
2613     }
2614 
2615     WTSFreeMemory(sessions);
2616     WTSFreeMemory(buffer_user);
2617     WTSFreeMemory(buffer_addr);
2618     return py_retlist;
2619 
2620 error:
2621     Py_XDECREF(py_username);
2622     Py_XDECREF(py_tuple);
2623     Py_XDECREF(py_address);
2624     Py_DECREF(py_retlist);
2625 
2626     if (sessions != NULL)
2627         WTSFreeMemory(sessions);
2628     if (buffer_user != NULL)
2629         WTSFreeMemory(buffer_user);
2630     if (buffer_addr != NULL)
2631         WTSFreeMemory(buffer_addr);
2632     return NULL;
2633 }
2634 
2635 
2636 /*
2637  * Return the number of handles opened by process.
2638  */
2639 static PyObject *
psutil_proc_num_handles(PyObject * self,PyObject * args)2640 psutil_proc_num_handles(PyObject *self, PyObject *args) {
2641     DWORD pid;
2642     HANDLE hProcess;
2643     DWORD handleCount;
2644 
2645     if (! PyArg_ParseTuple(args, "l", &pid))
2646         return NULL;
2647     hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION);
2648     if (NULL == hProcess)
2649         return NULL;
2650     if (! GetProcessHandleCount(hProcess, &handleCount)) {
2651         PyErr_SetFromWindowsErr(0);
2652         CloseHandle(hProcess);
2653         return NULL;
2654     }
2655     CloseHandle(hProcess);
2656     return Py_BuildValue("k", handleCount);
2657 }
2658 
2659 
2660 /*
2661  * Get various process information by using NtQuerySystemInformation.
2662  * We use this as a fallback when faster functions fail with access
2663  * denied. This is slower because it iterates over all processes.
2664  * Returned tuple includes the following process info:
2665  *
2666  * - num_threads()
2667  * - ctx_switches()
2668  * - num_handles() (fallback)
2669  * - cpu_times() (fallback)
2670  * - create_time() (fallback)
2671  * - io_counters() (fallback)
2672  * - memory_info() (fallback)
2673  */
2674 static PyObject *
psutil_proc_info(PyObject * self,PyObject * args)2675 psutil_proc_info(PyObject *self, PyObject *args) {
2676     DWORD pid;
2677     PSYSTEM_PROCESS_INFORMATION process;
2678     PVOID buffer;
2679     ULONG i;
2680     ULONG ctx_switches = 0;
2681     double user_time;
2682     double kernel_time;
2683     long long create_time;
2684     SIZE_T mem_private;
2685     PyObject *py_retlist;
2686 
2687     if (! PyArg_ParseTuple(args, "l", &pid))
2688         return NULL;
2689     if (! psutil_get_proc_info(pid, &process, &buffer))
2690         return NULL;
2691 
2692     for (i = 0; i < process->NumberOfThreads; i++)
2693         ctx_switches += process->Threads[i].ContextSwitches;
2694     user_time = (double)process->UserTime.HighPart * HI_T + \
2695                 (double)process->UserTime.LowPart * LO_T;
2696     kernel_time = (double)process->KernelTime.HighPart * HI_T + \
2697                     (double)process->KernelTime.LowPart * LO_T;
2698 
2699     // Convert the LARGE_INTEGER union to a Unix time.
2700     // It's the best I could find by googling and borrowing code here
2701     // and there. The time returned has a precision of 1 second.
2702     if (0 == pid || 4 == pid) {
2703         // the python module will translate this into BOOT_TIME later
2704         create_time = 0;
2705     }
2706     else {
2707         create_time = ((LONGLONG)process->CreateTime.HighPart) << 32;
2708         create_time += process->CreateTime.LowPart - 116444736000000000LL;
2709         create_time /= 10000000;
2710     }
2711 
2712 #if (_WIN32_WINNT >= 0x0501)  // Windows XP with SP2
2713     mem_private = process->PrivatePageCount;
2714 #else
2715     mem_private = 0;
2716 #endif
2717 
2718     py_retlist = Py_BuildValue(
2719 #if defined(_WIN64)
2720         "kkdddiKKKKKK" "kKKKKKKKKK",
2721 #else
2722         "kkdddiKKKKKK" "kIIIIIIIII",
2723 #endif
2724         process->HandleCount,                   // num handles
2725         ctx_switches,                           // num ctx switches
2726         user_time,                              // cpu user time
2727         kernel_time,                            // cpu kernel time
2728         (double)create_time,                    // create time
2729         (int)process->NumberOfThreads,          // num threads
2730         // IO counters
2731         process->ReadOperationCount.QuadPart,   // io rcount
2732         process->WriteOperationCount.QuadPart,  // io wcount
2733         process->ReadTransferCount.QuadPart,    // io rbytes
2734         process->WriteTransferCount.QuadPart,   // io wbytes
2735         process->OtherOperationCount.QuadPart,  // io others count
2736         process->OtherTransferCount.QuadPart,   // io others bytes
2737         // memory
2738         process->PageFaultCount,                // num page faults
2739         process->PeakWorkingSetSize,            // peak wset
2740         process->WorkingSetSize,                // wset
2741         process->QuotaPeakPagedPoolUsage,       // peak paged pool
2742         process->QuotaPagedPoolUsage,           // paged pool
2743         process->QuotaPeakNonPagedPoolUsage,    // peak non paged pool
2744         process->QuotaNonPagedPoolUsage,        // non paged pool
2745         process->PagefileUsage,                 // pagefile
2746         process->PeakPagefileUsage,             // peak pagefile
2747         mem_private                             // private
2748     );
2749 
2750     free(buffer);
2751     return py_retlist;
2752 }
2753 
2754 
get_region_protection_string(ULONG protection)2755 static char *get_region_protection_string(ULONG protection) {
2756     switch (protection & 0xff) {
2757         case PAGE_NOACCESS:
2758             return "";
2759         case PAGE_READONLY:
2760             return "r";
2761         case PAGE_READWRITE:
2762             return "rw";
2763         case PAGE_WRITECOPY:
2764             return "wc";
2765         case PAGE_EXECUTE:
2766             return "x";
2767         case PAGE_EXECUTE_READ:
2768             return "xr";
2769         case PAGE_EXECUTE_READWRITE:
2770             return "xrw";
2771         case PAGE_EXECUTE_WRITECOPY:
2772             return "xwc";
2773         default:
2774             return "?";
2775     }
2776 }
2777 
2778 
2779 /*
2780  * Return a list of process's memory mappings.
2781  */
2782 static PyObject *
psutil_proc_memory_maps(PyObject * self,PyObject * args)2783 psutil_proc_memory_maps(PyObject *self, PyObject *args) {
2784     MEMORY_BASIC_INFORMATION basicInfo;
2785     DWORD pid;
2786     HANDLE hProcess = NULL;
2787     PVOID baseAddress;
2788     ULONGLONG previousAllocationBase;
2789     WCHAR mappedFileName[MAX_PATH];
2790     LPVOID maxAddr;
2791     // required by GetMappedFileNameW
2792     DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
2793     PyObject *py_retlist = PyList_New(0);
2794     PyObject *py_tuple = NULL;
2795     PyObject *py_str = NULL;
2796 
2797     if (py_retlist == NULL)
2798         return NULL;
2799     if (! PyArg_ParseTuple(args, "l", &pid))
2800         goto error;
2801     hProcess = psutil_handle_from_pid(pid, access);
2802     if (NULL == hProcess)
2803         goto error;
2804 
2805     maxAddr = PSUTIL_SYSTEM_INFO.lpMaximumApplicationAddress;
2806     baseAddress = NULL;
2807 
2808     while (VirtualQueryEx(hProcess, baseAddress, &basicInfo,
2809                           sizeof(MEMORY_BASIC_INFORMATION)))
2810     {
2811         py_tuple = NULL;
2812         if (baseAddress > maxAddr)
2813             break;
2814         if (GetMappedFileNameW(hProcess, baseAddress, mappedFileName,
2815                                sizeof(mappedFileName)))
2816         {
2817             py_str = PyUnicode_FromWideChar(mappedFileName,
2818                                             wcslen(mappedFileName));
2819             if (py_str == NULL)
2820                 goto error;
2821 #ifdef _WIN64
2822            py_tuple = Py_BuildValue(
2823               "(KsOI)",
2824               (unsigned long long)baseAddress,
2825 #else
2826            py_tuple = Py_BuildValue(
2827               "(ksOI)",
2828               (unsigned long)baseAddress,
2829 #endif
2830               get_region_protection_string(basicInfo.Protect),
2831               py_str,
2832               basicInfo.RegionSize);
2833 
2834             if (!py_tuple)
2835                 goto error;
2836             if (PyList_Append(py_retlist, py_tuple))
2837                 goto error;
2838             Py_CLEAR(py_tuple);
2839             Py_CLEAR(py_str);
2840         }
2841         previousAllocationBase = (ULONGLONG)basicInfo.AllocationBase;
2842         baseAddress = (PCHAR)baseAddress + basicInfo.RegionSize;
2843     }
2844 
2845     CloseHandle(hProcess);
2846     return py_retlist;
2847 
2848 error:
2849     Py_XDECREF(py_tuple);
2850     Py_XDECREF(py_str);
2851     Py_DECREF(py_retlist);
2852     if (hProcess != NULL)
2853         CloseHandle(hProcess);
2854     return NULL;
2855 }
2856 
2857 
2858 /*
2859  * Return a {pid:ppid, ...} dict for all running processes.
2860  */
2861 static PyObject *
psutil_ppid_map(PyObject * self,PyObject * args)2862 psutil_ppid_map(PyObject *self, PyObject *args) {
2863     PyObject *py_pid = NULL;
2864     PyObject *py_ppid = NULL;
2865     PyObject *py_retdict = PyDict_New();
2866     HANDLE handle = NULL;
2867     PROCESSENTRY32 pe = {0};
2868     pe.dwSize = sizeof(PROCESSENTRY32);
2869 
2870     if (py_retdict == NULL)
2871         return NULL;
2872     handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
2873     if (handle == INVALID_HANDLE_VALUE) {
2874         PyErr_SetFromWindowsErr(0);
2875         Py_DECREF(py_retdict);
2876         return NULL;
2877     }
2878 
2879     if (Process32First(handle, &pe)) {
2880         do {
2881             py_pid = Py_BuildValue("I", pe.th32ProcessID);
2882             if (py_pid == NULL)
2883                 goto error;
2884             py_ppid = Py_BuildValue("I", pe.th32ParentProcessID);
2885             if (py_ppid == NULL)
2886                 goto error;
2887             if (PyDict_SetItem(py_retdict, py_pid, py_ppid))
2888                 goto error;
2889             Py_CLEAR(py_pid);
2890             Py_CLEAR(py_ppid);
2891         } while (Process32Next(handle, &pe));
2892     }
2893 
2894     CloseHandle(handle);
2895     return py_retdict;
2896 
2897 error:
2898     Py_XDECREF(py_pid);
2899     Py_XDECREF(py_ppid);
2900     Py_DECREF(py_retdict);
2901     CloseHandle(handle);
2902     return NULL;
2903 }
2904 
2905 
2906 /*
2907  * Return NICs addresses.
2908  */
2909 
2910 static PyObject *
psutil_net_if_addrs(PyObject * self,PyObject * args)2911 psutil_net_if_addrs(PyObject *self, PyObject *args) {
2912     unsigned int i = 0;
2913     ULONG family;
2914     PCTSTR intRet;
2915     PCTSTR netmaskIntRet;
2916     char *ptr;
2917     char buff_addr[1024];
2918     char buff_macaddr[1024];
2919     char buff_netmask[1024];
2920     DWORD dwRetVal = 0;
2921 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
2922     ULONG converted_netmask;
2923     UINT netmask_bits;
2924     struct in_addr in_netmask;
2925 #endif
2926     PIP_ADAPTER_ADDRESSES pAddresses = NULL;
2927     PIP_ADAPTER_ADDRESSES pCurrAddresses = NULL;
2928     PIP_ADAPTER_UNICAST_ADDRESS pUnicast = NULL;
2929 
2930     PyObject *py_retlist = PyList_New(0);
2931     PyObject *py_tuple = NULL;
2932     PyObject *py_address = NULL;
2933     PyObject *py_mac_address = NULL;
2934     PyObject *py_nic_name = NULL;
2935     PyObject *py_netmask = NULL;
2936 
2937     if (py_retlist == NULL)
2938         return NULL;
2939 
2940     pAddresses = psutil_get_nic_addresses();
2941     if (pAddresses == NULL)
2942         goto error;
2943     pCurrAddresses = pAddresses;
2944 
2945     while (pCurrAddresses) {
2946         pUnicast = pCurrAddresses->FirstUnicastAddress;
2947 
2948         netmaskIntRet = NULL;
2949         py_nic_name = NULL;
2950         py_nic_name = PyUnicode_FromWideChar(
2951             pCurrAddresses->FriendlyName,
2952             wcslen(pCurrAddresses->FriendlyName));
2953         if (py_nic_name == NULL)
2954             goto error;
2955 
2956         // MAC address
2957         if (pCurrAddresses->PhysicalAddressLength != 0) {
2958             ptr = buff_macaddr;
2959             *ptr = '\0';
2960             for (i = 0; i < (int) pCurrAddresses->PhysicalAddressLength; i++) {
2961                 if (i == (pCurrAddresses->PhysicalAddressLength - 1)) {
2962                     sprintf_s(ptr, _countof(buff_macaddr), "%.2X\n",
2963                             (int)pCurrAddresses->PhysicalAddress[i]);
2964                 }
2965                 else {
2966                     sprintf_s(ptr, _countof(buff_macaddr), "%.2X-",
2967                             (int)pCurrAddresses->PhysicalAddress[i]);
2968                 }
2969                 ptr += 3;
2970             }
2971             *--ptr = '\0';
2972 
2973             py_mac_address = Py_BuildValue("s", buff_macaddr);
2974             if (py_mac_address == NULL)
2975                 goto error;
2976 
2977             Py_INCREF(Py_None);
2978             Py_INCREF(Py_None);
2979             Py_INCREF(Py_None);
2980             py_tuple = Py_BuildValue(
2981                 "(OiOOOO)",
2982                 py_nic_name,
2983                 -1,  // this will be converted later to AF_LINK
2984                 py_mac_address,
2985                 Py_None,  // netmask (not supported)
2986                 Py_None,  // broadcast (not supported)
2987                 Py_None  // ptp (not supported on Windows)
2988             );
2989             if (! py_tuple)
2990                 goto error;
2991             if (PyList_Append(py_retlist, py_tuple))
2992                 goto error;
2993             Py_CLEAR(py_tuple);
2994             Py_CLEAR(py_mac_address);
2995         }
2996 
2997         // find out the IP address associated with the NIC
2998         if (pUnicast != NULL) {
2999             for (i = 0; pUnicast != NULL; i++) {
3000                 family = pUnicast->Address.lpSockaddr->sa_family;
3001                 if (family == AF_INET) {
3002                     struct sockaddr_in *sa_in = (struct sockaddr_in *)
3003                         pUnicast->Address.lpSockaddr;
3004                     intRet = inet_ntop(AF_INET, &(sa_in->sin_addr), buff_addr,
3005                                        sizeof(buff_addr));
3006                     if (!intRet)
3007                         goto error;
3008 #if (_WIN32_WINNT >= 0x0600) // Windows Vista and above
3009                     netmask_bits = pUnicast->OnLinkPrefixLength;
3010                     dwRetVal = ConvertLengthToIpv4Mask(netmask_bits, &converted_netmask);
3011                     if (dwRetVal == NO_ERROR) {
3012                         in_netmask.s_addr = converted_netmask;
3013                         netmaskIntRet = inet_ntop(
3014                             AF_INET, &in_netmask, buff_netmask,
3015                             sizeof(buff_netmask));
3016                         if (!netmaskIntRet)
3017                             goto error;
3018                     }
3019 #endif
3020                 }
3021                 else if (family == AF_INET6) {
3022                     struct sockaddr_in6 *sa_in6 = (struct sockaddr_in6 *)
3023                         pUnicast->Address.lpSockaddr;
3024                     intRet = inet_ntop(AF_INET6, &(sa_in6->sin6_addr),
3025                                        buff_addr, sizeof(buff_addr));
3026                     if (!intRet)
3027                         goto error;
3028                 }
3029                 else {
3030                     // we should never get here
3031                     pUnicast = pUnicast->Next;
3032                     continue;
3033                 }
3034 
3035 #if PY_MAJOR_VERSION >= 3
3036                 py_address = PyUnicode_FromString(buff_addr);
3037 #else
3038                 py_address = PyString_FromString(buff_addr);
3039 #endif
3040                 if (py_address == NULL)
3041                     goto error;
3042 
3043                 if (netmaskIntRet != NULL) {
3044 #if PY_MAJOR_VERSION >= 3
3045                     py_netmask = PyUnicode_FromString(buff_netmask);
3046 #else
3047                     py_netmask = PyString_FromString(buff_netmask);
3048 #endif
3049                 } else {
3050                     Py_INCREF(Py_None);
3051                     py_netmask = Py_None;
3052                 }
3053 
3054                 Py_INCREF(Py_None);
3055                 Py_INCREF(Py_None);
3056                 py_tuple = Py_BuildValue(
3057                     "(OiOOOO)",
3058                     py_nic_name,
3059                     family,
3060                     py_address,
3061                     py_netmask,
3062                     Py_None,  // broadcast (not supported)
3063                     Py_None  // ptp (not supported on Windows)
3064                 );
3065 
3066                 if (! py_tuple)
3067                     goto error;
3068                 if (PyList_Append(py_retlist, py_tuple))
3069                     goto error;
3070                 Py_CLEAR(py_tuple);
3071                 Py_CLEAR(py_address);
3072                 Py_CLEAR(py_netmask);
3073 
3074                 pUnicast = pUnicast->Next;
3075             }
3076         }
3077         Py_CLEAR(py_nic_name);
3078         pCurrAddresses = pCurrAddresses->Next;
3079     }
3080 
3081     free(pAddresses);
3082     return py_retlist;
3083 
3084 error:
3085     if (pAddresses)
3086         free(pAddresses);
3087     Py_DECREF(py_retlist);
3088     Py_XDECREF(py_tuple);
3089     Py_XDECREF(py_address);
3090     Py_XDECREF(py_nic_name);
3091     Py_XDECREF(py_netmask);
3092     return NULL;
3093 }
3094 
3095 
3096 /*
3097  * Provides stats about NIC interfaces installed on the system.
3098  * TODO: get 'duplex' (currently it's hard coded to '2', aka
3099          'full duplex')
3100  */
3101 static PyObject *
psutil_net_if_stats(PyObject * self,PyObject * args)3102 psutil_net_if_stats(PyObject *self, PyObject *args) {
3103     int i;
3104     DWORD dwSize = 0;
3105     DWORD dwRetVal = 0;
3106     MIB_IFTABLE *pIfTable;
3107     MIB_IFROW *pIfRow;
3108     PIP_ADAPTER_ADDRESSES pAddresses = NULL;
3109     PIP_ADAPTER_ADDRESSES pCurrAddresses = NULL;
3110     char descr[MAX_PATH];
3111     int ifname_found;
3112 
3113     PyObject *py_nic_name = NULL;
3114     PyObject *py_retdict = PyDict_New();
3115     PyObject *py_ifc_info = NULL;
3116     PyObject *py_is_up = NULL;
3117 
3118     if (py_retdict == NULL)
3119         return NULL;
3120 
3121     pAddresses = psutil_get_nic_addresses();
3122     if (pAddresses == NULL)
3123         goto error;
3124 
3125     pIfTable = (MIB_IFTABLE *) malloc(sizeof (MIB_IFTABLE));
3126     if (pIfTable == NULL) {
3127         PyErr_NoMemory();
3128         goto error;
3129     }
3130     dwSize = sizeof(MIB_IFTABLE);
3131     if (GetIfTable(pIfTable, &dwSize, FALSE) == ERROR_INSUFFICIENT_BUFFER) {
3132         free(pIfTable);
3133         pIfTable = (MIB_IFTABLE *) malloc(dwSize);
3134         if (pIfTable == NULL) {
3135             PyErr_NoMemory();
3136             goto error;
3137         }
3138     }
3139     // Make a second call to GetIfTable to get the actual
3140     // data we want.
3141     if ((dwRetVal = GetIfTable(pIfTable, &dwSize, FALSE)) != NO_ERROR) {
3142         PyErr_SetString(PyExc_RuntimeError, "GetIfTable() syscall failed");
3143         goto error;
3144     }
3145 
3146     for (i = 0; i < (int) pIfTable->dwNumEntries; i++) {
3147         pIfRow = (MIB_IFROW *) & pIfTable->table[i];
3148 
3149         // GetIfTable is not able to give us NIC with "friendly names"
3150         // so we determine them via GetAdapterAddresses() which
3151         // provides friendly names *and* descriptions and find the
3152         // ones that match.
3153         ifname_found = 0;
3154         pCurrAddresses = pAddresses;
3155         while (pCurrAddresses) {
3156             sprintf_s(descr, MAX_PATH, "%wS", pCurrAddresses->Description);
3157             if (lstrcmp(descr, pIfRow->bDescr) == 0) {
3158                 py_nic_name = PyUnicode_FromWideChar(
3159                     pCurrAddresses->FriendlyName,
3160                     wcslen(pCurrAddresses->FriendlyName));
3161                 if (py_nic_name == NULL)
3162                     goto error;
3163                 ifname_found = 1;
3164                 break;
3165             }
3166             pCurrAddresses = pCurrAddresses->Next;
3167         }
3168         if (ifname_found == 0) {
3169             // Name not found means GetAdapterAddresses() doesn't list
3170             // this NIC, only GetIfTable, meaning it's not really a NIC
3171             // interface so we skip it.
3172             continue;
3173         }
3174 
3175         // is up?
3176         if((pIfRow->dwOperStatus == MIB_IF_OPER_STATUS_CONNECTED ||
3177                 pIfRow->dwOperStatus == MIB_IF_OPER_STATUS_OPERATIONAL) &&
3178                 pIfRow->dwAdminStatus == 1 ) {
3179             py_is_up = Py_True;
3180         }
3181         else {
3182             py_is_up = Py_False;
3183         }
3184         Py_INCREF(py_is_up);
3185 
3186         py_ifc_info = Py_BuildValue(
3187             "(Oikk)",
3188             py_is_up,
3189             2,  // there's no way to know duplex so let's assume 'full'
3190             pIfRow->dwSpeed / 1000000,  // expressed in bytes, we want Mb
3191             pIfRow->dwMtu
3192         );
3193         if (!py_ifc_info)
3194             goto error;
3195         if (PyDict_SetItem(py_retdict, py_nic_name, py_ifc_info))
3196             goto error;
3197         Py_CLEAR(py_nic_name);
3198         Py_CLEAR(py_ifc_info);
3199     }
3200 
3201     free(pIfTable);
3202     free(pAddresses);
3203     return py_retdict;
3204 
3205 error:
3206     Py_XDECREF(py_is_up);
3207     Py_XDECREF(py_ifc_info);
3208     Py_XDECREF(py_nic_name);
3209     Py_DECREF(py_retdict);
3210     if (pIfTable != NULL)
3211         free(pIfTable);
3212     if (pAddresses != NULL)
3213         free(pAddresses);
3214     return NULL;
3215 }
3216 
3217 
3218 /*
3219  * Return CPU statistics.
3220  */
3221 static PyObject *
psutil_cpu_stats(PyObject * self,PyObject * args)3222 psutil_cpu_stats(PyObject *self, PyObject *args) {
3223     NTSTATUS status;
3224     _SYSTEM_PERFORMANCE_INFORMATION *spi = NULL;
3225     _SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *sppi = NULL;
3226     _SYSTEM_INTERRUPT_INFORMATION *InterruptInformation = NULL;
3227     unsigned int ncpus;
3228     UINT i;
3229     ULONG64 dpcs = 0;
3230     ULONG interrupts = 0;
3231 
3232     // retrieves number of processors
3233     ncpus = psutil_get_num_cpus(1);
3234     if (ncpus == 0)
3235         goto error;
3236 
3237     // get syscalls / ctx switches
3238     spi = (_SYSTEM_PERFORMANCE_INFORMATION *) \
3239            malloc(ncpus * sizeof(_SYSTEM_PERFORMANCE_INFORMATION));
3240     if (spi == NULL) {
3241         PyErr_NoMemory();
3242         goto error;
3243     }
3244     status = psutil_NtQuerySystemInformation(
3245         SystemPerformanceInformation,
3246         spi,
3247         ncpus * sizeof(_SYSTEM_PERFORMANCE_INFORMATION),
3248         NULL);
3249     if (! NT_SUCCESS(status)) {
3250         psutil_SetFromNTStatusErr(
3251             status, "NtQuerySystemInformation(SystemPerformanceInformation)");
3252         goto error;
3253     }
3254 
3255     // get DPCs
3256     InterruptInformation = \
3257         malloc(sizeof(_SYSTEM_INTERRUPT_INFORMATION) * ncpus);
3258     if (InterruptInformation == NULL) {
3259         PyErr_NoMemory();
3260         goto error;
3261     }
3262 
3263     status = psutil_NtQuerySystemInformation(
3264         SystemInterruptInformation,
3265         InterruptInformation,
3266         ncpus * sizeof(SYSTEM_INTERRUPT_INFORMATION),
3267         NULL);
3268     if (! NT_SUCCESS(status)) {
3269         psutil_SetFromNTStatusErr(
3270             status, "NtQuerySystemInformation(SystemInterruptInformation)");
3271         goto error;
3272     }
3273     for (i = 0; i < ncpus; i++) {
3274         dpcs += InterruptInformation[i].DpcCount;
3275     }
3276 
3277     // get interrupts
3278     sppi = (_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *) \
3279         malloc(ncpus * sizeof(_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION));
3280     if (sppi == NULL) {
3281         PyErr_NoMemory();
3282         goto error;
3283     }
3284 
3285     status = psutil_NtQuerySystemInformation(
3286         SystemProcessorPerformanceInformation,
3287         sppi,
3288         ncpus * sizeof(_SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION),
3289         NULL);
3290     if (! NT_SUCCESS(status)) {
3291         psutil_SetFromNTStatusErr(
3292             status,
3293             "NtQuerySystemInformation(SystemProcessorPerformanceInformation)");
3294         goto error;
3295     }
3296 
3297     for (i = 0; i < ncpus; i++) {
3298         interrupts += sppi[i].InterruptCount;
3299     }
3300 
3301     // done
3302     free(spi);
3303     free(InterruptInformation);
3304     free(sppi);
3305     return Py_BuildValue(
3306         "kkkk",
3307         spi->ContextSwitches,
3308         interrupts,
3309         (unsigned long)dpcs,
3310         spi->SystemCalls
3311     );
3312 
3313 error:
3314     if (spi)
3315         free(spi);
3316     if (InterruptInformation)
3317         free(InterruptInformation);
3318     if (sppi)
3319         free(sppi);
3320     return NULL;
3321 }
3322 
3323 
3324 /*
3325  * Return CPU frequency.
3326  */
3327 static PyObject *
psutil_cpu_freq(PyObject * self,PyObject * args)3328 psutil_cpu_freq(PyObject *self, PyObject *args) {
3329     PROCESSOR_POWER_INFORMATION *ppi;
3330     NTSTATUS ret;
3331     ULONG size;
3332     LPBYTE pBuffer = NULL;
3333     ULONG current;
3334     ULONG max;
3335     unsigned int ncpus;
3336 
3337     // Get the number of CPUs.
3338     ncpus = psutil_get_num_cpus(1);
3339     if (ncpus == 0)
3340         return NULL;
3341 
3342     // Allocate size.
3343     size = ncpus * sizeof(PROCESSOR_POWER_INFORMATION);
3344     pBuffer = (BYTE*)LocalAlloc(LPTR, size);
3345     if (! pBuffer)
3346         return PyErr_SetFromWindowsErr(0);
3347 
3348     // Syscall.
3349     ret = CallNtPowerInformation(
3350         ProcessorInformation, NULL, 0, pBuffer, size);
3351     if (ret != 0) {
3352         PyErr_SetString(PyExc_RuntimeError,
3353                         "CallNtPowerInformation syscall failed");
3354         goto error;
3355     }
3356 
3357     // Results.
3358     ppi = (PROCESSOR_POWER_INFORMATION *)pBuffer;
3359     max = ppi->MaxMhz;
3360     current = ppi->CurrentMhz;
3361     LocalFree(pBuffer);
3362 
3363     return Py_BuildValue("kk", current, max);
3364 
3365 error:
3366     if (pBuffer != NULL)
3367         LocalFree(pBuffer);
3368     return NULL;
3369 }
3370 
3371 
3372 /*
3373  * Return battery usage stats.
3374  */
3375 static PyObject *
psutil_sensors_battery(PyObject * self,PyObject * args)3376 psutil_sensors_battery(PyObject *self, PyObject *args) {
3377     SYSTEM_POWER_STATUS sps;
3378 
3379     if (GetSystemPowerStatus(&sps) == 0)
3380         return PyErr_SetFromWindowsErr(0);
3381     return Py_BuildValue(
3382         "iiiI",
3383         sps.ACLineStatus,  // whether AC is connected: 0=no, 1=yes, 255=unknown
3384         // status flag:
3385         // 1, 2, 4 = high, low, critical
3386         // 8 = charging
3387         // 128 = no battery
3388         sps.BatteryFlag,
3389         sps.BatteryLifePercent,  // percent
3390         sps.BatteryLifeTime  // remaining secs
3391     );
3392 }
3393 
3394 
3395 /*
3396  * System memory page size as an int.
3397  */
3398 static PyObject *
psutil_getpagesize(PyObject * self,PyObject * args)3399 psutil_getpagesize(PyObject *self, PyObject *args) {
3400     // XXX: we may want to use GetNativeSystemInfo to differentiate
3401     // page size for WoW64 processes (but am not sure).
3402     return Py_BuildValue("I", PSUTIL_SYSTEM_INFO.dwPageSize);
3403 }
3404 
3405 
3406 // ------------------------ Python init ---------------------------
3407 
3408 static PyMethodDef
3409 PsutilMethods[] = {
3410     // --- per-process functions
3411     {"proc_cmdline", (PyCFunction)(void(*)(void))psutil_proc_cmdline,
3412         METH_VARARGS | METH_KEYWORDS,
3413      "Return process cmdline as a list of cmdline arguments"},
3414     {"proc_environ", psutil_proc_environ, METH_VARARGS,
3415      "Return process environment data"},
3416     {"proc_exe", psutil_proc_exe, METH_VARARGS,
3417      "Return path of the process executable"},
3418     {"proc_name", psutil_proc_name, METH_VARARGS,
3419      "Return process name"},
3420     {"proc_kill", psutil_proc_kill, METH_VARARGS,
3421      "Kill the process identified by the given PID"},
3422     {"proc_cpu_times", psutil_proc_cpu_times, METH_VARARGS,
3423      "Return tuple of user/kern time for the given PID"},
3424     {"proc_create_time", psutil_proc_create_time, METH_VARARGS,
3425      "Return a float indicating the process create time expressed in "
3426      "seconds since the epoch"},
3427     {"proc_memory_info", psutil_proc_memory_info, METH_VARARGS,
3428      "Return a tuple of process memory information"},
3429     {"proc_memory_uss", psutil_proc_memory_uss, METH_VARARGS,
3430      "Return the USS of the process"},
3431     {"proc_cwd", psutil_proc_cwd, METH_VARARGS,
3432      "Return process current working directory"},
3433     {"proc_suspend_or_resume", psutil_proc_suspend_or_resume, METH_VARARGS,
3434      "Suspend or resume a process"},
3435     {"proc_open_files", psutil_proc_open_files, METH_VARARGS,
3436      "Return files opened by process"},
3437     {"proc_username", psutil_proc_username, METH_VARARGS,
3438      "Return the username of a process"},
3439     {"proc_threads", psutil_proc_threads, METH_VARARGS,
3440      "Return process threads information as a list of tuple"},
3441     {"proc_wait", psutil_proc_wait, METH_VARARGS,
3442      "Wait for process to terminate and return its exit code."},
3443     {"proc_priority_get", psutil_proc_priority_get, METH_VARARGS,
3444      "Return process priority."},
3445     {"proc_priority_set", psutil_proc_priority_set, METH_VARARGS,
3446      "Set process priority."},
3447 #if (_WIN32_WINNT >= 0x0600)  // Windows Vista
3448     {"proc_io_priority_get", psutil_proc_io_priority_get, METH_VARARGS,
3449      "Return process IO priority."},
3450     {"proc_io_priority_set", psutil_proc_io_priority_set, METH_VARARGS,
3451      "Set process IO priority."},
3452 #endif
3453     {"proc_cpu_affinity_get", psutil_proc_cpu_affinity_get, METH_VARARGS,
3454      "Return process CPU affinity as a bitmask."},
3455     {"proc_cpu_affinity_set", psutil_proc_cpu_affinity_set, METH_VARARGS,
3456      "Set process CPU affinity."},
3457     {"proc_io_counters", psutil_proc_io_counters, METH_VARARGS,
3458      "Get process I/O counters."},
3459     {"proc_is_suspended", psutil_proc_is_suspended, METH_VARARGS,
3460      "Return True if one of the process threads is in a suspended state"},
3461     {"proc_num_handles", psutil_proc_num_handles, METH_VARARGS,
3462      "Return the number of handles opened by process."},
3463     {"proc_memory_maps", psutil_proc_memory_maps, METH_VARARGS,
3464      "Return a list of process's memory mappings"},
3465 
3466     // --- alternative pinfo interface
3467     {"proc_info", psutil_proc_info, METH_VARARGS,
3468      "Various process information"},
3469 
3470     // --- system-related functions
3471     {"pids", psutil_pids, METH_VARARGS,
3472      "Returns a list of PIDs currently running on the system"},
3473     {"ppid_map", psutil_ppid_map, METH_VARARGS,
3474      "Return a {pid:ppid, ...} dict for all running processes"},
3475     {"pid_exists", psutil_pid_exists, METH_VARARGS,
3476      "Determine if the process exists in the current process list."},
3477     {"cpu_count_logical", psutil_cpu_count_logical, METH_VARARGS,
3478      "Returns the number of logical CPUs on the system"},
3479     {"cpu_count_phys", psutil_cpu_count_phys, METH_VARARGS,
3480      "Returns the number of physical CPUs on the system"},
3481     {"boot_time", psutil_boot_time, METH_VARARGS,
3482      "Return the system boot time expressed in seconds since the epoch."},
3483     {"virtual_mem", psutil_virtual_mem, METH_VARARGS,
3484      "Return the total amount of physical memory, in bytes"},
3485     {"cpu_times", psutil_cpu_times, METH_VARARGS,
3486      "Return system cpu times as a list"},
3487     {"per_cpu_times", psutil_per_cpu_times, METH_VARARGS,
3488      "Return system per-cpu times as a list of tuples"},
3489     {"disk_usage", psutil_disk_usage, METH_VARARGS,
3490      "Return path's disk total and free as a Python tuple."},
3491     {"net_io_counters", psutil_net_io_counters, METH_VARARGS,
3492      "Return dict of tuples of networks I/O information."},
3493     {"disk_io_counters", psutil_disk_io_counters, METH_VARARGS,
3494      "Return dict of tuples of disks I/O information."},
3495     {"users", psutil_users, METH_VARARGS,
3496      "Return a list of currently connected users."},
3497     {"disk_partitions", psutil_disk_partitions, METH_VARARGS,
3498      "Return disk partitions."},
3499     {"net_connections", psutil_net_connections, METH_VARARGS,
3500      "Return system-wide connections"},
3501     {"net_if_addrs", psutil_net_if_addrs, METH_VARARGS,
3502      "Return NICs addresses."},
3503     {"net_if_stats", psutil_net_if_stats, METH_VARARGS,
3504      "Return NICs stats."},
3505     {"cpu_stats", psutil_cpu_stats, METH_VARARGS,
3506      "Return NICs stats."},
3507     {"cpu_freq", psutil_cpu_freq, METH_VARARGS,
3508      "Return CPU frequency."},
3509 #if (_WIN32_WINNT >= 0x0600)  // Windows Vista
3510     {"init_loadavg_counter", (PyCFunction)psutil_init_loadavg_counter,
3511      METH_VARARGS,
3512      "Initializes the emulated load average calculator."},
3513     {"getloadavg", (PyCFunction)psutil_get_loadavg, METH_VARARGS,
3514      "Returns the emulated POSIX-like load average."},
3515 #endif
3516     {"sensors_battery", psutil_sensors_battery, METH_VARARGS,
3517      "Return battery metrics usage."},
3518     {"getpagesize", psutil_getpagesize, METH_VARARGS,
3519      "Return system memory page size."},
3520 
3521     // --- windows services
3522     {"winservice_enumerate", psutil_winservice_enumerate, METH_VARARGS,
3523      "List all services"},
3524     {"winservice_query_config", psutil_winservice_query_config, METH_VARARGS,
3525      "Return service config"},
3526     {"winservice_query_status", psutil_winservice_query_status, METH_VARARGS,
3527      "Return service config"},
3528     {"winservice_query_descr", psutil_winservice_query_descr, METH_VARARGS,
3529      "Return the description of a service"},
3530     {"winservice_start", psutil_winservice_start, METH_VARARGS,
3531      "Start a service"},
3532     {"winservice_stop", psutil_winservice_stop, METH_VARARGS,
3533      "Stop a service"},
3534 
3535     // --- windows API bindings
3536     {"win32_QueryDosDevice", psutil_win32_QueryDosDevice, METH_VARARGS,
3537      "QueryDosDevice binding"},
3538 
3539     // --- others
3540     {"set_testing", psutil_set_testing, METH_NOARGS,
3541      "Set psutil in testing mode"},
3542 
3543     {NULL, NULL, 0, NULL}
3544 };
3545 
3546 
3547 struct module_state {
3548     PyObject *error;
3549 };
3550 
3551 #if PY_MAJOR_VERSION >= 3
3552 #define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
3553 #else
3554 #define GETSTATE(m) (&_state)
3555 static struct module_state _state;
3556 #endif
3557 
3558 #if PY_MAJOR_VERSION >= 3
3559 
psutil_windows_traverse(PyObject * m,visitproc visit,void * arg)3560 static int psutil_windows_traverse(PyObject *m, visitproc visit, void *arg) {
3561     Py_VISIT(GETSTATE(m)->error);
3562     return 0;
3563 }
3564 
psutil_windows_clear(PyObject * m)3565 static int psutil_windows_clear(PyObject *m) {
3566     Py_CLEAR(GETSTATE(m)->error);
3567     return 0;
3568 }
3569 
3570 static struct PyModuleDef moduledef = {
3571     PyModuleDef_HEAD_INIT,
3572     "psutil_windows",
3573     NULL,
3574     sizeof(struct module_state),
3575     PsutilMethods,
3576     NULL,
3577     psutil_windows_traverse,
3578     psutil_windows_clear,
3579     NULL
3580 };
3581 
3582 #define INITERROR return NULL
3583 
PyInit__psutil_windows(void)3584 PyMODINIT_FUNC PyInit__psutil_windows(void)
3585 
3586 #else
3587 #define INITERROR return
3588 void init_psutil_windows(void)
3589 #endif
3590 {
3591     struct module_state *st = NULL;
3592 #if PY_MAJOR_VERSION >= 3
3593     PyObject *module = PyModule_Create(&moduledef);
3594 #else
3595     PyObject *module = Py_InitModule("_psutil_windows", PsutilMethods);
3596 #endif
3597     if (module == NULL)
3598         INITERROR;
3599 
3600     if (psutil_setup() != 0)
3601         INITERROR;
3602     if (psutil_load_globals() != 0)
3603         INITERROR;
3604     if (psutil_set_se_debug() != 0)
3605         INITERROR;
3606 
3607     st = GETSTATE(module);
3608     st->error = PyErr_NewException("_psutil_windows.Error", NULL, NULL);
3609     if (st->error == NULL) {
3610         Py_DECREF(module);
3611         INITERROR;
3612     }
3613 
3614     // Exceptions.
3615     TimeoutExpired = PyErr_NewException(
3616         "_psutil_windows.TimeoutExpired", NULL, NULL);
3617     Py_INCREF(TimeoutExpired);
3618     PyModule_AddObject(module, "TimeoutExpired", TimeoutExpired);
3619 
3620     TimeoutAbandoned = PyErr_NewException(
3621         "_psutil_windows.TimeoutAbandoned", NULL, NULL);
3622     Py_INCREF(TimeoutAbandoned);
3623     PyModule_AddObject(module, "TimeoutAbandoned", TimeoutAbandoned);
3624 
3625     // version constant
3626     PyModule_AddIntConstant(module, "version", PSUTIL_VERSION);
3627 
3628     // process status constants
3629     // http://msdn.microsoft.com/en-us/library/ms683211(v=vs.85).aspx
3630     PyModule_AddIntConstant(
3631         module, "ABOVE_NORMAL_PRIORITY_CLASS", ABOVE_NORMAL_PRIORITY_CLASS);
3632     PyModule_AddIntConstant(
3633         module, "BELOW_NORMAL_PRIORITY_CLASS", BELOW_NORMAL_PRIORITY_CLASS);
3634     PyModule_AddIntConstant(
3635         module, "HIGH_PRIORITY_CLASS", HIGH_PRIORITY_CLASS);
3636     PyModule_AddIntConstant(
3637         module, "IDLE_PRIORITY_CLASS", IDLE_PRIORITY_CLASS);
3638     PyModule_AddIntConstant(
3639         module, "NORMAL_PRIORITY_CLASS", NORMAL_PRIORITY_CLASS);
3640     PyModule_AddIntConstant(
3641         module, "REALTIME_PRIORITY_CLASS", REALTIME_PRIORITY_CLASS);
3642 
3643     // connection status constants
3644     // http://msdn.microsoft.com/en-us/library/cc669305.aspx
3645     PyModule_AddIntConstant(
3646         module, "MIB_TCP_STATE_CLOSED", MIB_TCP_STATE_CLOSED);
3647     PyModule_AddIntConstant(
3648         module, "MIB_TCP_STATE_CLOSING", MIB_TCP_STATE_CLOSING);
3649     PyModule_AddIntConstant(
3650         module, "MIB_TCP_STATE_CLOSE_WAIT", MIB_TCP_STATE_CLOSE_WAIT);
3651     PyModule_AddIntConstant(
3652         module, "MIB_TCP_STATE_LISTEN", MIB_TCP_STATE_LISTEN);
3653     PyModule_AddIntConstant(
3654         module, "MIB_TCP_STATE_ESTAB", MIB_TCP_STATE_ESTAB);
3655     PyModule_AddIntConstant(
3656         module, "MIB_TCP_STATE_SYN_SENT", MIB_TCP_STATE_SYN_SENT);
3657     PyModule_AddIntConstant(
3658         module, "MIB_TCP_STATE_SYN_RCVD", MIB_TCP_STATE_SYN_RCVD);
3659     PyModule_AddIntConstant(
3660         module, "MIB_TCP_STATE_FIN_WAIT1", MIB_TCP_STATE_FIN_WAIT1);
3661     PyModule_AddIntConstant(
3662         module, "MIB_TCP_STATE_FIN_WAIT2", MIB_TCP_STATE_FIN_WAIT2);
3663     PyModule_AddIntConstant(
3664         module, "MIB_TCP_STATE_LAST_ACK", MIB_TCP_STATE_LAST_ACK);
3665     PyModule_AddIntConstant(
3666         module, "MIB_TCP_STATE_TIME_WAIT", MIB_TCP_STATE_TIME_WAIT);
3667     PyModule_AddIntConstant(
3668         module, "MIB_TCP_STATE_TIME_WAIT", MIB_TCP_STATE_TIME_WAIT);
3669     PyModule_AddIntConstant(
3670         module, "MIB_TCP_STATE_DELETE_TCB", MIB_TCP_STATE_DELETE_TCB);
3671     PyModule_AddIntConstant(
3672         module, "PSUTIL_CONN_NONE", PSUTIL_CONN_NONE);
3673 
3674     // service status constants
3675     /*
3676     PyModule_AddIntConstant(
3677         module, "SERVICE_CONTINUE_PENDING", SERVICE_CONTINUE_PENDING);
3678     PyModule_AddIntConstant(
3679         module, "SERVICE_PAUSE_PENDING", SERVICE_PAUSE_PENDING);
3680     PyModule_AddIntConstant(
3681         module, "SERVICE_PAUSED", SERVICE_PAUSED);
3682     PyModule_AddIntConstant(
3683         module, "SERVICE_RUNNING", SERVICE_RUNNING);
3684     PyModule_AddIntConstant(
3685         module, "SERVICE_START_PENDING", SERVICE_START_PENDING);
3686     PyModule_AddIntConstant(
3687         module, "SERVICE_STOP_PENDING", SERVICE_STOP_PENDING);
3688     PyModule_AddIntConstant(
3689         module, "SERVICE_STOPPED", SERVICE_STOPPED);
3690     */
3691 
3692     // ...for internal use in _psutil_windows.py
3693     PyModule_AddIntConstant(
3694         module, "INFINITE", INFINITE);
3695     PyModule_AddIntConstant(
3696         module, "ERROR_ACCESS_DENIED", ERROR_ACCESS_DENIED);
3697     PyModule_AddIntConstant(
3698         module, "ERROR_INVALID_NAME", ERROR_INVALID_NAME);
3699     PyModule_AddIntConstant(
3700         module, "ERROR_SERVICE_DOES_NOT_EXIST", ERROR_SERVICE_DOES_NOT_EXIST);
3701     PyModule_AddIntConstant(
3702         module, "ERROR_PRIVILEGE_NOT_HELD", ERROR_PRIVILEGE_NOT_HELD);
3703     PyModule_AddIntConstant(
3704         module, "WINVER", PSUTIL_WINVER);
3705     PyModule_AddIntConstant(
3706         module, "WINDOWS_XP", PSUTIL_WINDOWS_XP);
3707     PyModule_AddIntConstant(
3708         module, "WINDOWS_SERVER_2003", PSUTIL_WINDOWS_SERVER_2003);
3709     PyModule_AddIntConstant(
3710         module, "WINDOWS_VISTA", PSUTIL_WINDOWS_VISTA);
3711     PyModule_AddIntConstant(
3712         module, "WINDOWS_7", PSUTIL_WINDOWS_7);
3713     PyModule_AddIntConstant(
3714         module, "WINDOWS_8", PSUTIL_WINDOWS_8);
3715     PyModule_AddIntConstant(
3716         module, "WINDOWS_8_1", PSUTIL_WINDOWS_8_1);
3717     PyModule_AddIntConstant(
3718         module, "WINDOWS_10", PSUTIL_WINDOWS_10);
3719 
3720 #if PY_MAJOR_VERSION >= 3
3721     return module;
3722 #endif
3723 }
3724