1  #include "pch.h"
  2  
  3  #include "FileLocksmith.h"
  4  #include "NtdllExtensions.h"
  5  
  6  static bool is_directory(const std::wstring path)
  7  {
  8      DWORD attributes = GetFileAttributesW(path.c_str());
  9      return attributes != INVALID_FILE_ATTRIBUTES && attributes & FILE_ATTRIBUTE_DIRECTORY;
 10  }
 11  
 12  // C++20 method
 13  static constexpr bool starts_with(std::wstring_view whole, std::wstring_view part)
 14  {
 15      return whole.size() >= part.size() && whole.substr(0, part.size()) == part;
 16  }
 17  
 18  std::vector<ProcessResult> find_processes_recursive(const std::vector<std::wstring>& paths)
 19  {
 20      NtdllExtensions nt_ext;
 21  
 22      // TODO use a trie!
 23  
 24      // This maps kernel names of files within `paths` to their normal paths.
 25      std::map<std::wstring, std::wstring> kernel_names_files;
 26  
 27      // This maps kernel names of directories within `paths` to their normal paths.
 28      std::map<std::wstring, std::wstring> kernel_names_dirs;
 29  
 30      for (const auto& path : paths)
 31      {
 32          auto kernel_path = nt_ext.path_to_kernel_name(path.c_str());
 33          if (!kernel_path.empty())
 34          {
 35              (is_directory(path) ? kernel_names_dirs : kernel_names_files)[kernel_path] = path;
 36          }
 37      }
 38  
 39      std::map<ULONG_PTR, std::set<std::wstring>> pid_files;
 40  
 41      // Returns a normal path of the file specified by kernel_name, if it matches
 42      // the search criteria. Otherwise, return an empty string.
 43      auto kernel_paths_contain = [&](const std::wstring& kernel_name) -> std::wstring
 44      {
 45          // Normal equivalence
 46          if (auto it = kernel_names_files.find(kernel_name); it != kernel_names_files.end())
 47          {
 48              return it->second;
 49          }
 50  
 51          if (auto it = kernel_names_dirs.find(kernel_name); it != kernel_names_dirs.end())
 52          {
 53              return it->second;
 54          }
 55  
 56          for (const auto& [dir_kernel_name, dir_path] : kernel_names_dirs)
 57          {
 58              if (starts_with(kernel_name, dir_kernel_name + (dir_kernel_name.length()>0&&dir_kernel_name[dir_kernel_name.length()-1]!=L'\\' ? L"\\" : L"")))
 59              {
 60                  return dir_path + kernel_name.substr(dir_kernel_name.size());
 61              }
 62          }
 63  
 64          return {};
 65      };
 66  
 67      for (const auto& handle_info : nt_ext.handles())
 68      {
 69          if (handle_info.type_name == L"File")
 70          {
 71              auto path = kernel_paths_contain(handle_info.kernel_file_name);
 72              if (!path.empty())
 73              {
 74                  pid_files[handle_info.pid].insert(std::move(path));
 75              }
 76          }
 77      }
 78  
 79      // Check all modules used by processes
 80      auto processes = nt_ext.processes();
 81  
 82      for (const auto& process : processes)
 83      {
 84          for (const auto& path : process.modules)
 85          {
 86              auto kernel_name = nt_ext.path_to_kernel_name(path.c_str());
 87  
 88              auto found_path = kernel_paths_contain(kernel_name);
 89              if (!found_path.empty())
 90              {
 91                  pid_files[process.pid].insert(std::move(found_path));
 92              }
 93          }
 94      }
 95  
 96      std::vector<ProcessResult> result;
 97  
 98      for (const auto& process_info : processes)
 99      {
100          if (auto it = pid_files.find(process_info.pid); it != pid_files.end())
101          {
102              result.push_back(ProcessResult
103                  {
104                      process_info.name,
105                      process_info.pid,
106                      process_info.user,
107                      std::vector(it->second.begin(), it->second.end())
108                  });
109          }
110      }
111  
112      return result;
113  }
114  
115  constexpr size_t LongMaxPathSize = 65536;
116  
117  std::wstring pid_to_full_path(DWORD pid)
118  {
119      HANDLE process = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
120  
121      std::wstring result(LongMaxPathSize, L'\0');
122  
123      // Returns zero on failure, so it's okay to resize to zero.
124      auto length = GetModuleFileNameExW(process, NULL, result.data(), static_cast<DWORD>(result.size()));
125      result.resize(length);
126  
127      CloseHandle(process);
128      return result;
129  }