Process SE DEBUG mode was not set on Windows (#1429)

6 files changed, 179 insertions, 173 deletions

diff --git a/Makefile b/Makefile
index 25928e0a..74480518 100644
--- a/Makefile
+++ b/Makefile
@@ -266,5 +266,9 @@ bench-oneshot-2:  ## Same as above but using perf module (supposed to be more pr
 check-broken-links:  ## Look for broken links in source files.
 	git ls-files | xargs $(PYTHON) -Wa scripts/internal/check_broken_links.py
 
+print-access-denied:
+#	${MAKE} install
+	$(TEST_PREFIX) $(PYTHON) scripts/internal/procs_access_denied.py
+
 help: ## Display callable targets.
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}'
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 7dc37ee1..111cd71a 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -3724,7 +3724,9 @@ void init_psutil_windows(void)
         module, "ERROR_SERVICE_DOES_NOT_EXIST", ERROR_SERVICE_DOES_NOT_EXIST);
 
     // set SeDebug for the current process
-    psutil_set_se_debug();
+    if (psutil_set_se_debug() != 0)
+        return NULL;
+
     psutil_setup();
     if (psutil_load_globals() != 0)
         return NULL;
diff --git a/psutil/arch/windows/security.c b/psutil/arch/windows/security.c
index d5f8f8d3..4e2c7435 100644
--- a/psutil/arch/windows/security.c
+++ b/psutil/arch/windows/security.c
@@ -4,134 +4,44 @@
  * found in the LICENSE file.
  *
  * Security related functions for Windows platform (Set privileges such as
- * SeDebug), as well as security helper functions.
+ * SE DEBUG).
  */
 
 #include <windows.h>
 #include <Python.h>
+#include "../../_psutil_common.h"
 
 
-/*
- * Convert a process handle to a process token handle.
- */
-HANDLE
-psutil_token_from_handle(HANDLE hProcess) {
-    HANDLE hToken = NULL;
-
-    if (! OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
-        return PyErr_SetFromWindowsErr(0);
-    return hToken;
-}
-
-
-/*
- * http://www.ddj.com/windows/184405986
- *
- * There's a way to determine whether we're running under the Local System
- * account. However (you guessed it), we have to call more Win32 functions to
- * determine this. Backing up through the code listing, we need to make another
- * call to GetTokenInformation, but instead of passing through the TOKEN_USER
- * constant, we pass through the TOKEN_PRIVILEGES constant. This value returns
- * an array of privileges that the account has in the environment. Iterating
- * through the array, we call the function LookupPrivilegeName looking for the
- * string “SeTcbPrivilege. If the function returns this string, then this
- * account has Local System privileges
- */
-int
-psutil_has_system_privilege(HANDLE hProcess) {
-    DWORD i;
-    DWORD dwSize = 0;
-    DWORD dwRetval = 0;
-    TCHAR privName[256];
-    DWORD dwNameSize = 256;
-    // PTOKEN_PRIVILEGES tp = NULL;
-    BYTE *pBuffer = NULL;
-    TOKEN_PRIVILEGES *tp = NULL;
-    HANDLE hToken = psutil_token_from_handle(hProcess);
-
-    if (NULL == hToken)
-        return -1;
-    // call GetTokenInformation first to get the buffer size
-    if (! GetTokenInformation(hToken, TokenPrivileges, NULL, 0, &dwSize)) {
-        dwRetval = GetLastError();
-        // if it failed for a reason other than the buffer, bail out
-        if (dwRetval != ERROR_INSUFFICIENT_BUFFER ) {
-            PyErr_SetFromWindowsErr(dwRetval);
-            return 0;
-        }
-    }
-
-    // allocate buffer and call GetTokenInformation again
-    // tp = (PTOKEN_PRIVILEGES) GlobalAlloc(GPTR, dwSize);
-    pBuffer = (BYTE *) malloc(dwSize);
-    if (pBuffer == NULL) {
-        PyErr_NoMemory();
-        return -1;
-    }
-
-    if (! GetTokenInformation(hToken, TokenPrivileges, pBuffer,
-                              dwSize, &dwSize))
-    {
-        PyErr_SetFromWindowsErr(0);
-        free(pBuffer);
-        return -1;
-    }
-
-    // convert the BYTE buffer to a TOKEN_PRIVILEGES struct pointer
-    tp = (TOKEN_PRIVILEGES *)pBuffer;
-
-    // check all the privileges looking for SeTcbPrivilege
-    for (i = 0; i < tp->PrivilegeCount; i++) {
-        // reset the buffer contents and the buffer size
-        strcpy(privName, "");
-        dwNameSize = sizeof(privName) / sizeof(TCHAR);
-        if (! LookupPrivilegeName(NULL,
-                                  &tp->Privileges[i].Luid,
-                                  (LPTSTR)privName,
-                                  &dwNameSize))
-        {
-            PyErr_SetFromWindowsErr(0);
-            free(pBuffer);
-            return -1;
-        }
-
-        // if we find the SeTcbPrivilege then it's a LocalSystem process
-        if (! lstrcmpi(privName, TEXT("SeTcbPrivilege"))) {
-            free(pBuffer);
-            return 1;
-        }
-    }
-
-    free(pBuffer);
-    return 0;
-}
-
-
-BOOL
+static BOOL
 psutil_set_privilege(HANDLE hToken, LPCTSTR Privilege, BOOL bEnablePrivilege) {
     TOKEN_PRIVILEGES tp;
     LUID luid;
     TOKEN_PRIVILEGES tpPrevious;
     DWORD cbPrevious = sizeof(TOKEN_PRIVILEGES);
 
-    if (!LookupPrivilegeValue( NULL, Privilege, &luid )) return FALSE;
+    if (! LookupPrivilegeValue(NULL, Privilege, &luid)) {
+        PyErr_SetFromOSErrnoWithSyscall("LookupPrivilegeValue");
+        return 1;
+    }
 
     // first pass.  get current privilege setting
     tp.PrivilegeCount = 1;
     tp.Privileges[0].Luid = luid;
     tp.Privileges[0].Attributes = 0;
 
-    AdjustTokenPrivileges(
-        hToken,
-        FALSE,
-        &tp,
-        sizeof(TOKEN_PRIVILEGES),
-        &tpPrevious,
-        &cbPrevious
-    );
-
-    if (GetLastError() != ERROR_SUCCESS) return FALSE;
-    // second pass. set privilege based on previous setting
+    if (! AdjustTokenPrivileges(
+            hToken,
+            FALSE,
+            &tp,
+            sizeof(TOKEN_PRIVILEGES),
+            &tpPrevious,
+            &cbPrevious))
+    {
+        PyErr_SetFromOSErrnoWithSyscall("AdjustTokenPrivileges");
+        return 1;
+    }
+
+    // Second pass. Set privilege based on previous setting.
     tpPrevious.PrivilegeCount = 1;
     tpPrevious.Privileges[0].Luid = luid;
 
@@ -141,81 +51,88 @@ psutil_set_privilege(HANDLE hToken, LPCTSTR Privilege, BOOL bEnablePrivilege) {
         tpPrevious.Privileges[0].Attributes ^=
             (SE_PRIVILEGE_ENABLED & tpPrevious.Privileges[0].Attributes);
 
-    AdjustTokenPrivileges(
-        hToken,
-        FALSE,
-        &tpPrevious,
-        cbPrevious,
-        NULL,
-        NULL
-    );
-
-    if (GetLastError() != ERROR_SUCCESS) return FALSE;
+    if (! AdjustTokenPrivileges(
+            hToken,
+            FALSE,
+            &tpPrevious,
+            cbPrevious,
+            NULL,
+            NULL))
+    {
+        PyErr_SetFromOSErrnoWithSyscall("AdjustTokenPrivileges");
+        return 1;
+    }
 
-    return TRUE;
+    return 0;
 }
 
 
-int
-psutil_set_se_debug() {
-    HANDLE hToken;
-    if (!OpenProcessToken(GetCurrentProcess(),
-                          TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
-                          &hToken)
-       ) {
-        if (GetLastError() == ERROR_NO_TOKEN) {
-            if (!ImpersonateSelf(SecurityImpersonation)) {
-                CloseHandle(hToken);
-                return 0;
+static HANDLE
+psutil_get_thisproc_token() {
+    HANDLE hToken = NULL;
+    HANDLE me = GetCurrentProcess();
+
+    if (! OpenProcessToken(
+            me, TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
+    {
+        if (GetLastError() == ERROR_NO_TOKEN)
+        {
+            if (! ImpersonateSelf(SecurityImpersonation)) {
+                PyErr_SetFromOSErrnoWithSyscall("ImpersonateSelf");
+                return NULL;
             }
-            if (!OpenProcessToken(GetCurrentProcess(),
-                                 TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
-                                 &hToken)
-               ) {
-                RevertToSelf();
-                CloseHandle(hToken);
-                return 0;
+            if (! OpenProcessToken(
+                    me, TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
+            {
+                PyErr_SetFromOSErrnoWithSyscall("OpenProcessToken");
+                return NULL;
             }
         }
+        else {
+            PyErr_SetFromOSErrnoWithSyscall("OpenProcessToken");
+            return NULL;
+        }
     }
 
-    // enable SeDebugPrivilege (open any process)
-    if (! psutil_set_privilege(hToken, SE_DEBUG_NAME, TRUE)) {
-        RevertToSelf();
-        CloseHandle(hToken);
-        return 0;
-    }
+    return hToken;
+}
 
-    RevertToSelf();
-    CloseHandle(hToken);
-    return 1;
+
+static void
+psutil_print_err() {
+    char *msg = "psutil module couldn't set SE DEBUG mode for this process; " \
+        "please file an issue against psutil bug tracker";
+    psutil_debug(msg);
+    if (GetLastError() != ERROR_ACCESS_DENIED)
+        PyErr_WarnEx(PyExc_RuntimeWarning, msg, 1);
+    PyErr_Clear();
 }
 
 
+/*
+ * Set this process in SE DEBUG mode so that we have more chances of
+ * querying processes owned by other users, including many owned by
+ * Administrator and Local System.
+ * https://docs.microsoft.com/windows-hardware/drivers/debugger/debug-privilege
+ * This is executed on module import and we don't crash on error.
+ */
 int
-psutil_unset_se_debug() {
+psutil_set_se_debug() {
     HANDLE hToken;
-    if (!OpenProcessToken(GetCurrentProcess(),
-                          TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
-                          &hToken)
-       ) {
-        if (GetLastError() == ERROR_NO_TOKEN) {
-            if (! ImpersonateSelf(SecurityImpersonation))
-                return 0;
-            if (!OpenProcessToken(GetCurrentProcess(),
-                                 TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
-                                 &hToken))
-            {
-                return 0;
-            }
-        }
-    }
+    int err = 1;
 
-    // now disable SeDebug
-    if (! psutil_set_privilege(hToken, SE_DEBUG_NAME, FALSE))
+    if ((hToken = psutil_get_thisproc_token()) == NULL) {
+        // "return 1;" to get an exception
+        psutil_print_err();
         return 0;
+    }
+
+    if (psutil_set_privilege(hToken, SE_DEBUG_NAME, TRUE) != 0) {
+        // "return 1;" to get an exception
+        psutil_print_err();
+    }
 
+    RevertToSelf();
     CloseHandle(hToken);
-    return 1;
+    return 0;
 }
-
diff --git a/psutil/arch/windows/security.h b/psutil/arch/windows/security.h
index aa8a22ad..8d4ddb00 100644
--- a/psutil/arch/windows/security.h
+++ b/psutil/arch/windows/security.h
@@ -9,9 +9,5 @@
 
 #include <windows.h>
 
-BOOL psutil_set_privilege(HANDLE hToken, LPCTSTR Privilege, BOOL bEnablePrivilege);
-HANDLE psutil_token_from_handle(HANDLE hProcess);
-int psutil_has_system_privilege(HANDLE hProcess);
 int psutil_set_se_debug();
-int psutil_unset_se_debug();
 
diff --git a/scripts/internal/procs_access_denied.py b/scripts/internal/procs_access_denied.py
new file mode 100644
index 00000000..9f792480
--- /dev/null
+++ b/scripts/internal/procs_access_denied.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""
+Helper script which tries to access all info of all running processes.
+It prints how many AccessDenied exceptions are raised in total and
+for each Process method.
+"""
+
+from __future__ import print_function, division
+from collections import defaultdict
+import sys
+
+import psutil
+
+
+def term_supports_colors(file=sys.stdout):
+    try:
+        import curses
+        assert file.isatty()
+        curses.setupterm()
+        assert curses.tigetnum("colors") > 0
+    except Exception:
+        return False
+    else:
+        return True
+
+
+COLORS = term_supports_colors()
+
+
+def hilite(s, ok=True, bold=False):
+    """Return an highlighted version of 'string'."""
+    if not COLORS:
+        return s
+    attr = []
+    if ok is None:  # no color
+        pass
+    elif ok:   # green
+        attr.append('32')
+    else:   # red
+        attr.append('31')
+    if bold:
+        attr.append('1')
+    return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
+
+
+def main():
+    tot_procs = 0
+    tot_ads = 0
+    signaler = object()
+    d = defaultdict(int)
+    for p in psutil.process_iter(attrs=[], ad_value=signaler):
+        tot_procs += 1
+        for methname, value in p.info.items():
+            if value is signaler:
+                tot_ads += 1
+                d[methname] += 1
+            else:
+                d[methname] += 0
+
+    for methname, ads in sorted(d.items(), key=lambda x: x[1]):
+        perc = (ads / tot_procs) * 100
+        s = "%-20s %-3s %5.1f%%   " % (methname, ads, perc)
+        if not ads:
+            s += "SUCCESS"
+            s = hilite(s, ok=True)
+        else:
+            s += "ACCESS DENIED"
+            s = hilite(s, ok=False)
+        print(s)
+    print("--------------------------")
+    print("total: %19s (%s total processes)" % (tot_ads, tot_procs))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/scripts/internal/winmake.py b/scripts/internal/winmake.py
index b1ce7b8a..49aae699 100755
--- a/scripts/internal/winmake.py
+++ b/scripts/internal/winmake.py
@@ -483,6 +483,13 @@ def bench_oneshot_2():
     sh("%s -Wa scripts\\internal\\bench_oneshot_2.py" % PYTHON)
 
 
+@cmd
+def print_access_denied():
+    """Benchmarks for oneshot() ctx manager (see #799)."""
+    install()
+    sh("%s -Wa scripts\\internal\\procs_access_denied.py" % PYTHON)
+
+
 def set_python(s):
     global PYTHON
     if os.path.isabs(s):