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