From 98f9263f805d5a84ad15a2f5086a2b8b4770ab26 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Wed, 14 Jun 2023 13:11:17 +0530
Subject: [PATCH] Add a utility function to get the process path for whichever
 process has a file open on Windows

---
 setup/extensions.json                 |  2 +-
 src/calibre/utils/windows/winutil.cpp | 54 +++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/setup/extensions.json b/setup/extensions.json
index cf7d145b66..2648ce85c7 100644
--- a/setup/extensions.json
+++ b/setup/extensions.json
@@ -173,7 +173,7 @@
         "only": "windows",
         "headers": "calibre/utils/cpp_binding.h calibre/utils/windows/common.h",
         "sources": "calibre/utils/windows/winutil.cpp",
-        "libraries": "shell32 wininet advapi32 gdi32",
+        "libraries": "shell32 wininet advapi32 gdi32 rstrtmgr",
         "cflags": "/X"
     },
     {
diff --git a/src/calibre/utils/windows/winutil.cpp b/src/calibre/utils/windows/winutil.cpp
index 7295a2886c..8aff12ed50 100644
--- a/src/calibre/utils/windows/winutil.cpp
+++ b/src/calibre/utils/windows/winutil.cpp
@@ -7,6 +7,7 @@
 
 #include "common.h"
 #include <processthreadsapi.h>
+#include <vector>
 #include <wininet.h>
 #include <lmcons.h>
 #include <combaseapi.h>
@@ -20,6 +21,7 @@
 #include <comdef.h>
 #include <atlbase.h>  // for CComPtr
 #include <versionhelpers.h>
+#include <restartmanager.h>
 
 // GUID {{{
 typedef struct {
@@ -470,6 +472,53 @@ winutil_read_directory_changes(PyObject *self, PyObject *args) {
     return ans;
 }
 
+static void rm_end_session(DWORD sid) { RmEndSession(sid); }
+
+static PyObject*
+winutil_get_processes_using_files(PyObject *self, PyObject *args) {
+    DWORD dwSession;
+    WCHAR szSessionKey[CCH_RM_SESSION_KEY+1] = { 0 };
+    DWORD dwError = RmStartSession(&dwSession, 0, szSessionKey);
+    if (dwError != ERROR_SUCCESS) { PyErr_SetFromWindowsErr(dwError); return NULL; }
+    generic_raii<DWORD, rm_end_session> sr(dwSession);
+    std::vector<wchar_raii> paths(PyTuple_GET_SIZE(args));
+    for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(args); i++) {
+        if (!py_to_wchar_no_none(PyTuple_GET_ITEM(args, i), &paths[i])) return NULL;
+    }
+    std::vector<const wchar_t*> array_of_paths(paths.size());
+    for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(args); i++) { array_of_paths[i] = paths[i].ptr(); }
+    dwError = RmRegisterResources(dwSession, array_of_paths.size(), array_of_paths.data(), 0, NULL, 0, NULL);
+    if (dwError != ERROR_SUCCESS) { PyErr_SetFromWindowsErr(dwError); return NULL; }
+    DWORD dwReason;
+    UINT nProcInfoNeeded = 64, nProcInfo;
+    std::vector<RM_PROCESS_INFO> rgpi;
+    do {
+        nProcInfo = 2*nProcInfoNeeded; nProcInfoNeeded = 0;
+        rgpi.resize(nProcInfo);
+        dwError = RmGetList(dwSession, &nProcInfoNeeded, &nProcInfo, rgpi.data(), &dwReason);
+    } while (dwError == ERROR_MORE_DATA);
+    if (dwError != ERROR_SUCCESS) { PyErr_SetFromWindowsErr(dwError); return NULL; }
+    pyobject_raii ans(PyList_New(0));
+    if (!ans) return NULL;
+    std::vector<wchar_t> process_path(MAX_PATH*16);
+    for (UINT i = 0; i < nProcInfo; i++) {
+        handle_raii_null process(OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, rgpi[i].Process.dwProcessId));
+        if (process) {
+            FILETIME ftCreate, ftExit, ftKernel, ftUser;
+            if (GetProcessTimes(process.ptr(), &ftCreate, &ftExit, &ftKernel, &ftUser) && CompareFileTime(
+                        &rgpi[i].Process.ProcessStartTime, &ftCreate) == 0) {
+                DWORD cch = process_path.size();
+                if (QueryFullProcessImageNameW(process.ptr(), 0, process_path.data(), &cch) && cch < process_path.size()) {
+                    pyobject_raii pp(Py_BuildValue("{su su# si}", "app_name", rgpi[i].strAppName, "path", process_path.data(), (Py_ssize_t)cch, "app_type", (int)rgpi[i].ApplicationType));
+                    if (!pp) return NULL;
+                    if (PyList_Append(ans.ptr(), pp.ptr()) != 0) return NULL;
+                }
+            }
+        }
+    }
+    return ans.detach();
+}
+
 static PyObject*
 winutil_get_file_size(PyObject *self, PyObject *args) {
 	HANDLE handle;
@@ -1302,6 +1351,11 @@ static PyMethodDef winutil_methods[] = {
         "read_directory_changes(handle, buffer, subtree, flags)\n\nWrapper for ReadDirectoryChangesW"
     },
 
+    {"get_processes_using_files", (PyCFunction)winutil_get_processes_using_files, METH_VARARGS,
+        "get_processes_using_files(path1, path2, ...)\n\nGet information about processes that have the specified files open."
+    },
+
+
     {NULL, NULL, 0, NULL}
 };
 #undef M