move windows cwd extraction code

Move cwd extraction code from _psutil_windows.c to process_info.c
so its closer to psutil_get_cmdline which accesses the process'
memory in a similar way.

3 files changed, 110 insertions, 97 deletions

diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 32f00b67..d8693fef 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -936,108 +936,18 @@ error:
 static PyObject *
 psutil_proc_cwd(PyObject *self, PyObject *args) {
     long pid;
-    HANDLE processHandle = NULL;
-    PVOID pebAddress;
-    PVOID rtlUserProcParamsAddress;
-    UNICODE_STRING currentDirectory;
-    WCHAR *currentDirectoryContent = NULL;
-    PyObject *py_unicode = NULL;
+    int pid_return;
 
     if (! PyArg_ParseTuple(args, "l", &pid))
         return NULL;
 
-    processHandle = psutil_handle_from_pid(pid);
-    if (processHandle == NULL)
+    pid_return = psutil_pid_is_running(pid);
+    if (pid_return == 0)
+        return NoSuchProcess();
+    if (pid_return == -1)
         return NULL;
 
-    pebAddress = psutil_get_peb_address(processHandle);
-
-    // get the address of ProcessParameters
-#ifdef _WIN64
-    if (!ReadProcessMemory(processHandle, (PCHAR)pebAddress + 32,
-                           &rtlUserProcParamsAddress, sizeof(PVOID), NULL))
-#else
-    if (!ReadProcessMemory(processHandle, (PCHAR)pebAddress + 0x10,
-                           &rtlUserProcParamsAddress, sizeof(PVOID), NULL))
-#endif
-    {
-        CloseHandle(processHandle);
-        if (GetLastError() == ERROR_PARTIAL_COPY) {
-            // this occurs quite often with system processes
-            return AccessDenied();
-        }
-        else {
-            return PyErr_SetFromWindowsErr(0);
-        }
-    }
-
-    // Read the currentDirectory UNICODE_STRING structure.
-    // 0x24 refers to "CurrentDirectoryPath" of RTL_USER_PROCESS_PARAMETERS
-    // structure, see:
-    // http://wj32.wordpress.com/2009/01/24/
-    //     howto-get-the-command-line-of-processes/
-#ifdef _WIN64
-    if (!ReadProcessMemory(processHandle, (PCHAR)rtlUserProcParamsAddress + 56,
-                           &currentDirectory, sizeof(currentDirectory), NULL))
-#else
-    if (!ReadProcessMemory(processHandle,
-                           (PCHAR)rtlUserProcParamsAddress + 0x24,
-                           &currentDirectory, sizeof(currentDirectory), NULL))
-#endif
-    {
-        CloseHandle(processHandle);
-        if (GetLastError() == ERROR_PARTIAL_COPY) {
-            // this occurs quite often with system processes
-            return AccessDenied();
-        }
-        else {
-            return PyErr_SetFromWindowsErr(0);
-        }
-    }
-
-    // allocate memory to hold cwd
-    currentDirectoryContent = (WCHAR *)malloc(currentDirectory.Length + 1);
-    if (currentDirectoryContent == NULL) {
-        PyErr_NoMemory();
-        goto error;
-    }
-
-    // read cwd
-    if (!ReadProcessMemory(processHandle, currentDirectory.Buffer,
-                           currentDirectoryContent, currentDirectory.Length,
-                           NULL))
-    {
-        if (GetLastError() == ERROR_PARTIAL_COPY) {
-            // this occurs quite often with system processes
-            AccessDenied();
-        }
-        else {
-            PyErr_SetFromWindowsErr(0);
-        }
-        goto error;
-    }
-
-    // null-terminate the string to prevent wcslen from returning
-    // incorrect length the length specifier is in characters, but
-    // currentDirectory.Length is in bytes
-    currentDirectoryContent[(currentDirectory.Length / sizeof(WCHAR))] = '\0';
-
-    // convert wchar array to a Python unicode string
-    py_unicode = PyUnicode_FromWideChar(
-        currentDirectoryContent, wcslen(currentDirectoryContent));
-    if (py_unicode == NULL)
-        goto error;
-    CloseHandle(processHandle);
-    free(currentDirectoryContent);
-    return py_unicode;
-
-error:
-    Py_XDECREF(py_unicode);
-    if (currentDirectoryContent != NULL)
-        free(currentDirectoryContent);
-    if (processHandle != NULL)
-        CloseHandle(processHandle);
-    return NULL;
+    return psutil_get_cwd(pid);
 }
 
 
diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c
index a064baa7..7273abb9 100644
--- a/psutil/arch/windows/process_info.c
+++ b/psutil/arch/windows/process_info.c
@@ -316,6 +316,109 @@ error:
 }
 
 
+PyObject *psutil_get_cwd(long pid) {
+    HANDLE processHandle = NULL;
+    PVOID pebAddress;
+    PVOID rtlUserProcParamsAddress;
+    UNICODE_STRING currentDirectory;
+    WCHAR *currentDirectoryContent = NULL;
+    PyObject *py_unicode = NULL;
+
+    processHandle = psutil_handle_from_pid(pid);
+    if (processHandle == NULL)
+        return NULL;
+
+    pebAddress = psutil_get_peb_address(processHandle);
+
+    // get the address of ProcessParameters
+#ifdef _WIN64
+    if (!ReadProcessMemory(processHandle, (PCHAR)pebAddress + 32,
+                           &rtlUserProcParamsAddress, sizeof(PVOID), NULL))
+#else
+    if (!ReadProcessMemory(processHandle, (PCHAR)pebAddress + 0x10,
+                           &rtlUserProcParamsAddress, sizeof(PVOID), NULL))
+#endif
+    {
+        CloseHandle(processHandle);
+        if (GetLastError() == ERROR_PARTIAL_COPY) {
+            // this occurs quite often with system processes
+            return AccessDenied();
+        }
+        else {
+            return PyErr_SetFromWindowsErr(0);
+        }
+    }
+
+    // Read the currentDirectory UNICODE_STRING structure.
+    // 0x24 refers to "CurrentDirectoryPath" of RTL_USER_PROCESS_PARAMETERS
+    // structure, see:
+    // http://wj32.wordpress.com/2009/01/24/
+    //     howto-get-the-command-line-of-processes/
+#ifdef _WIN64
+    if (!ReadProcessMemory(processHandle, (PCHAR)rtlUserProcParamsAddress + 56,
+                           &currentDirectory, sizeof(currentDirectory), NULL))
+#else
+    if (!ReadProcessMemory(processHandle,
+                           (PCHAR)rtlUserProcParamsAddress + 0x24,
+                           &currentDirectory, sizeof(currentDirectory), NULL))
+#endif
+    {
+        CloseHandle(processHandle);
+        if (GetLastError() == ERROR_PARTIAL_COPY) {
+            // this occurs quite often with system processes
+            return AccessDenied();
+        }
+        else {
+            return PyErr_SetFromWindowsErr(0);
+        }
+    }
+
+    // allocate memory to hold cwd
+    currentDirectoryContent = (WCHAR *)malloc(currentDirectory.Length + 1);
+    if (currentDirectoryContent == NULL) {
+        PyErr_NoMemory();
+        goto error;
+    }
+
+    // read cwd
+    if (!ReadProcessMemory(processHandle, currentDirectory.Buffer,
+                           currentDirectoryContent, currentDirectory.Length,
+                           NULL))
+    {
+        if (GetLastError() == ERROR_PARTIAL_COPY) {
+            // this occurs quite often with system processes
+            AccessDenied();
+        }
+        else {
+            PyErr_SetFromWindowsErr(0);
+        }
+        goto error;
+    }
+
+    // null-terminate the string to prevent wcslen from returning
+    // incorrect length the length specifier is in characters, but
+    // currentDirectory.Length is in bytes
+    currentDirectoryContent[(currentDirectory.Length / sizeof(WCHAR))] = '\0';
+
+    // convert wchar array to a Python unicode string
+    py_unicode = PyUnicode_FromWideChar(
+        currentDirectoryContent, wcslen(currentDirectoryContent));
+    if (py_unicode == NULL)
+        goto error;
+    CloseHandle(processHandle);
+    free(currentDirectoryContent);
+    return py_unicode;
+
+error:
+    Py_XDECREF(py_unicode);
+    if (currentDirectoryContent != NULL)
+        free(currentDirectoryContent);
+    if (processHandle != NULL)
+        CloseHandle(processHandle);
+    return NULL;
+}
+
+
 #define PH_FIRST_PROCESS(Processes) ((PSYSTEM_PROCESS_INFORMATION)(Processes))
 #define PH_NEXT_PROCESS(Process) ( \
    ((PSYSTEM_PROCESS_INFORMATION)(Process))->NextEntryOffset ? \
diff --git a/psutil/arch/windows/process_info.h b/psutil/arch/windows/process_info.h
index c2699192..908efc5f 100644
--- a/psutil/arch/windows/process_info.h
+++ b/psutil/arch/windows/process_info.h
@@ -18,8 +18,8 @@ HANDLE psutil_handle_from_pid_waccess(DWORD pid, DWORD dwDesiredAccess);
 int psutil_handlep_is_running(HANDLE hProcess);
 int psutil_pid_in_proclist(DWORD pid);
 int psutil_pid_is_running(DWORD pid);
-PVOID psutil_get_peb_address(HANDLE ProcessHandle);
 PyObject* psutil_get_cmdline(long pid);
+PyObject* psutil_get_cwd(long pid);
 int psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess,
                          PVOID *retBuffer);