From 631041d1b736630d57cbbe66593dd7ced1972171 Mon Sep 17 00:00:00 2001
From: Dan Ryan <dan@danryan.co>
Date: Thu, 8 Mar 2018 00:37:53 -0500
Subject: [PATCH] Remove dependence on PSUtil

- Use native ctypes for parent process traversal
- Addresses issues raised in https://github.com/pypa/pipenv/issues/1587
  and https://github.com/Microsoft/vscode-python/issues/978
- Improves speed on windows
- Allows pipenv to remove vendored psutil which sometimes fails to find
linked python dlls
---
 Pipfile           |   1 -
 pew/_win_utils.py | 107 ++++++++++++++++++++++++++++++++++++++++++++++
 pew/pew.py        |   4 +-
 requirements.txt  |   3 +-
 setup.py          |   3 --
 5 files changed, 110 insertions(+), 8 deletions(-)
 create mode 100644 pew/_win_utils.py

diff --git a/Pipfile b/Pipfile
index 8dfb404..720f73d 100644
--- a/Pipfile
+++ b/Pipfile
@@ -8,7 +8,6 @@ name = "pypi"
 virtualenv = "==1.11.6"
 "virtualenv-clone" = "==0.2.5"
 "pythonz-bd" = {"version" = "*", "sys_platform" = "!='win32'" }
-"psutil" = {"version" = "*", "sys_platform" = "=='win32'" }
 
 [dev-packages]
 
diff --git a/pew/_win_utils.py b/pew/_win_utils.py
new file mode 100644
index 0000000..b5afc0c
--- /dev/null
+++ b/pew/_win_utils.py
@@ -0,0 +1,107 @@
+# -*- coding=utf-8 -*-
+# psutil is painfully slow in win32. So to avoid adding big
+# dependencies like pywin32 a ctypes based solution is preferred
+
+# Code based on the winappdbg project http://winappdbg.sourceforge.net/
+# (BSD License) - adapted from Celery
+# https://github.com/celery/celery/blob/2.5-archived/celery/concurrency/processes/_win.py
+import os
+from ctypes import (
+    byref, sizeof, windll, Structure, WinError, POINTER,
+    c_size_t, c_char, c_void_p
+)
+from ctypes.wintypes import DWORD, LONG
+
+ERROR_NO_MORE_FILES = 18
+INVALID_HANDLE_VALUE = c_void_p(-1).value
+
+
+class PROCESSENTRY32(Structure):
+    _fields_ = [
+        ('dwSize', DWORD),
+        ('cntUsage', DWORD),
+        ('th32ProcessID', DWORD),
+        ('th32DefaultHeapID', c_size_t),
+        ('th32ModuleID', DWORD),
+        ('cntThreads', DWORD),
+        ('th32ParentProcessID', DWORD),
+        ('pcPriClassBase', LONG),
+        ('dwFlags', DWORD),
+        ('szExeFile', c_char * 260),
+    ]
+
+
+LPPROCESSENTRY32 = POINTER(PROCESSENTRY32)
+
+
+def CreateToolhelp32Snapshot(dwFlags=2, th32ProcessID=0):
+    hSnapshot = windll.kernel32.CreateToolhelp32Snapshot(
+        dwFlags,
+        th32ProcessID
+    )
+    if hSnapshot == INVALID_HANDLE_VALUE:
+        raise WinError()
+    return hSnapshot
+
+
+def Process32First(hSnapshot):
+    pe = PROCESSENTRY32()
+    pe.dwSize = sizeof(PROCESSENTRY32)
+    success = windll.kernel32.Process32First(hSnapshot, byref(pe))
+    if not success:
+        if windll.kernel32.GetLastError() == ERROR_NO_MORE_FILES:
+            return
+        raise WinError()
+    return pe
+
+
+def Process32Next(hSnapshot, pe=None):
+    if pe is None:
+        pe = PROCESSENTRY32()
+    pe.dwSize = sizeof(PROCESSENTRY32)
+    success = windll.kernel32.Process32Next(hSnapshot, byref(pe))
+    if not success:
+        if windll.kernel32.GetLastError() == ERROR_NO_MORE_FILES:
+            return
+        raise WinError()
+    return pe
+
+
+def get_all_processes():
+    """Return a dictionary of properties about all processes.
+
+    >>> get_all_processes()
+    {
+        1509: {
+            'parent_pid': 1201,
+            'executable': 'C:\\Program\\\\ Files\\Python36\\python.exe'
+        }
+    }
+    """
+    h_process = CreateToolhelp32Snapshot()
+    pids = {}
+    pe = Process32First(h_process)
+    while pe:
+        pids[pe.th32ProcessID] = {
+            'executable': str(pe.szExeFile.decode('utf-8')),
+        }
+        if pe.th32ParentProcessID:
+            pids[pe.th32ProcessID]['parent_pid'] = pe.th32ParentProcessID
+        pe = Process32Next(h_process, pe)
+
+    return pids
+
+
+def get_grandparent_process(pid=None):
+    """Get grandparent process name of the supplied pid or os.getpid().
+
+    :param int pid: The pid to track.
+    :return: Name of the grandparent process.
+    """
+    if not pid:
+        pid = os.getpid()
+    processes = get_all_processes()
+    ppid = processes[pid]['parent_pid']
+    parent = processes[ppid]
+    grandparent = processes[parent['parent_pid']]
+    return grandparent['executable']
diff --git a/pew/pew.py b/pew/pew.py
index bcaabad..470dfd3 100644
--- a/pew/pew.py
+++ b/pew/pew.py
@@ -36,7 +36,7 @@
     InstallCommand = ListPythons = LocatePython = UninstallCommand = \
         lambda : sys.exit('Command not supported on this platform')
 
-    import psutil
+    from ._win_utils import get_grandparent_process
 
 from pew._utils import (check_call, invoke, expandpath, own, env_bin_dir,
                         check_path, temp_environ, NamedTemporaryFile, to_unicode)
@@ -184,7 +184,7 @@ def _detect_shell():
         if 'CMDER_ROOT' in os.environ:
             shell = 'Cmder'
         elif windows:
-            shell = psutil.Process(os.getpid()).parent().parent().name()
+            shell = get_grandparent_process(os.getpid())
         else:
             shell = 'sh'
     return shell
diff --git a/requirements.txt b/requirements.txt
index 565f66d..f17ca6f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,5 +3,4 @@ virtualenv==1.11.6
 virtualenv-clone==0.2.5
 pytest==2.6.2
 pythonz-bd==1.11.2 ; sys_platform != 'win32'
-psutil==5.3.1 ; sys_platform == 'win32'
-stdeb ; sys_platform == 'linux'
\ No newline at end of file
+stdeb ; sys_platform == 'linux'
diff --git a/setup.py b/setup.py
index bc99fb5..76f966c 100644
--- a/setup.py
+++ b/setup.py
@@ -75,9 +75,6 @@ def run(self):
         ':python_version=="3.3"': [
             'pathlib'
         ],
-        ':sys_platform=="win32"': [
-            'psutil==5.3.1'
-        ],
         'pythonz': [
             'pythonz-bd>=1.10.2'
         ]