Mercurial > minori
diff dep/animia/src/win32.cpp @ 56:6ff7aabeb9d7
deps: add animia for open files detection
author | Paper <mrpapersonic@gmail.com> |
---|---|
date | Thu, 28 Sep 2023 12:35:21 -0400 |
parents | |
children | 4c6dd5999b39 |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/dep/animia/src/win32.cpp Thu Sep 28 12:35:21 2023 -0400 @@ -0,0 +1,258 @@ +#include "win32.h" +#include <windows.h> +#include <winternl.h> +#include <libloaderapi.h> +#include <ntdef.h> +#include <psapi.h> +#include <tlhelp32.h> +#include <fileapi.h> +#include <handleapi.h> +#include <vector> +#include <iostream> +#include <string> +#include <unordered_map> +#include <stdexcept> +#include <locale> +#include <codecvt> +/* This file is noticably more complex than Unix and Linux, and that's because + there is no "simple" way to get the paths of a file. */ + +#define SystemExtendedHandleInformation ((SYSTEM_INFORMATION_CLASS)0x40) +constexpr NTSTATUS STATUS_INFO_LENGTH_MISMATCH = 0xC0000004UL; +constexpr NTSTATUS STATUS_SUCCESS = 0x00000000UL; + +static unsigned short file_type_index = 0; + +struct SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX { + PVOID Object; + ULONG_PTR UniqueProcessId; + HANDLE HandleValue; + ACCESS_MASK GrantedAccess; + USHORT CreatorBackTraceIndex; + USHORT ObjectTypeIndex; + ULONG HandleAttributes; + ULONG Reserved; +}; + +struct SYSTEM_HANDLE_INFORMATION_EX { + ULONG_PTR NumberOfHandles; + ULONG_PTR Reserved; + SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1]; +}; + +namespace Animia::Windows { + +std::vector<int> get_all_pids() { + std::vector<int> ret; + HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); + PROCESSENTRY32 pe32; + pe32.dwSize = sizeof(PROCESSENTRY32); + + if (hProcessSnap == INVALID_HANDLE_VALUE) + return std::vector<int>(); + + if (!Process32First(hProcessSnap, &pe32)) + return std::vector<int>(); + + ret.push_back(pe32.th32ProcessID); + while (Process32Next(hProcessSnap, &pe32)) { + ret.push_back(pe32.th32ProcessID); + } + // clean the snapshot object + CloseHandle(hProcessSnap); + + return ret; +} + +std::string get_process_name(int pid) { + HANDLE handle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid); + if (!handle) + return ""; + + std::string ret(MAX_PATH, '\0'); + if (!GetModuleBaseNameA(handle, 0, &ret.front(), ret.size())) + throw std::runtime_error("GetModuleBaseNameA failed: " + std::to_string(GetLastError())); + CloseHandle(handle); + + return ret; +} + +/* All of this BS is required on Windows. Why? */ + +HANDLE DuplicateHandle(HANDLE process_handle, HANDLE handle) { + HANDLE dup_handle = nullptr; + const bool result = ::DuplicateHandle(process_handle, handle, + ::GetCurrentProcess(), &dup_handle, 0, false, DUPLICATE_SAME_ACCESS); + return result ? dup_handle : nullptr; +} + +PVOID GetNTDLLAddress(LPCSTR proc_name) { + return reinterpret_cast<PVOID>(GetProcAddress(GetModuleHandleA("ntdll.dll"), proc_name)); +} + +NTSTATUS QuerySystemInformation(SYSTEM_INFORMATION_CLASS cls, PVOID sysinfo, ULONG len, PULONG retlen) { + static const auto func = reinterpret_cast<decltype(::NtQuerySystemInformation)*>(GetNTDLLAddress("NtQuerySystemInformation")); + return func(cls, sysinfo, len, retlen); +} + +NTSTATUS QueryObject(HANDLE handle, OBJECT_INFORMATION_CLASS cls, PVOID objinf, ULONG objinflen, PULONG retlen) { + static const auto func = reinterpret_cast<decltype(::NtQueryObject)*>(GetNTDLLAddress("NtQueryObject")); + return func(handle, cls, objinf, objinflen, retlen); +} + +std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> GetSystemHandleInformation() { + std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> res; + ULONG cb = 1 << 19; + NTSTATUS status = STATUS_SUCCESS; + SYSTEM_HANDLE_INFORMATION_EX* info; + + do { + status = STATUS_NO_MEMORY; + + if (!(info = (SYSTEM_HANDLE_INFORMATION_EX*)malloc(cb *= 2))) + continue; + + if (0 <= (status = QuerySystemInformation(SystemExtendedHandleInformation, info, cb, &cb))) { + if (ULONG_PTR handles = info->NumberOfHandles) { + SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX* entry = info->Handles; + do { + if (entry) res.push_back(*entry); + } while (entry++, --handles); + } + } + free(info); + } while (status == STATUS_INFO_LENGTH_MISMATCH); + + return res; +} + +OBJECT_TYPE_INFORMATION QueryObjectTypeInfo(HANDLE handle) { + OBJECT_TYPE_INFORMATION info; + QueryObject(handle, ObjectTypeInformation, &info, sizeof(info), NULL); + return info; +} + +std::string UnicodeStringToStdString(UNICODE_STRING string) { + ANSI_STRING result; + static const auto uc_to_ansi = reinterpret_cast<decltype(::RtlUnicodeStringToAnsiString)*>(GetNTDLLAddress("RtlUnicodeStringToAnsiString")); + uc_to_ansi(&result, &string, TRUE); + std::string ret = std::string(result.Buffer, result.Length); + static const auto free_ansi = reinterpret_cast<decltype(::RtlFreeAnsiString)*>(GetNTDLLAddress("RtlFreeAnsiString")); + free_ansi(&result); + return ret; +} + +std::string GetHandleType(HANDLE handle) { + OBJECT_TYPE_INFORMATION info = QueryObjectTypeInfo(handle); + return UnicodeStringToStdString(info.TypeName); +} + +/* GetFinalPathNameByHandleA literally just doesn't work */ +std::string GetFinalPathNameByHandle(HANDLE handle) { + std::wstring buffer; + + int result = ::GetFinalPathNameByHandleW(handle, NULL, 0, FILE_NAME_NORMALIZED | VOLUME_NAME_DOS); + buffer.resize(result); + ::GetFinalPathNameByHandleW(handle, &buffer.front(), buffer.size(), FILE_NAME_NORMALIZED | VOLUME_NAME_DOS); + + std::wstring_convert<std::codecvt_utf8<wchar_t>, wchar_t> converter; + + return converter.to_bytes(buffer); +} + +std::string GetSystemDirectory() { + std::string windir = std::string(MAX_PATH, '\0'); + ::GetWindowsDirectoryA(&windir.front(), windir.length()); + return "\\\\?\\" + windir; +} + +/* This function is useless. I'm not exactly sure why, but whenever I try to compare the two + values, they both come up as different. I'm assuming it's just some Unicode BS I can't be bothered + to deal with. */ +bool IsSystemDirectory(const std::string& path) { + std::string path_l = path; + CharUpperBuffA(&path_l.front(), path_l.length()); + + std::string windir = GetSystemDirectory(); + CharUpperBuffA(&windir.front(), windir.length()); + + return path_l.rfind(windir, 0) != std::string::npos; +} + +bool IsFileHandle(HANDLE handle, unsigned short object_type_index) { + if (file_type_index) + return object_type_index == file_type_index; + else if (!handle) + return true; + else if (GetHandleType(handle) == "File") { + file_type_index = object_type_index; + return true; + } + return false; +} + +bool IsFileMaskOk(ACCESS_MASK access_mask) { + /* this filters out any file handles that, legitimately, + do not make sense (for what we're using it for) + + shoutout to erengy for having these in Anisthesia */ + + if (!(access_mask & FILE_READ_DATA)) + return false; + + if ((access_mask & FILE_APPEND_DATA) || + (access_mask & FILE_WRITE_EA) || + (access_mask & FILE_WRITE_ATTRIBUTES)) + return false; + + return true; +} + +bool IsFilePathOk(const std::string& path) { + if (path.empty() || IsSystemDirectory(path)) + return false; + + const auto file_attributes = GetFileAttributesA(path.c_str()); + if ((file_attributes == INVALID_FILE_ATTRIBUTES) || + (file_attributes & FILE_ATTRIBUTE_DIRECTORY)) + return false; + + return true; +} + +std::vector<std::string> get_open_files(int pid) { + std::unordered_map<int, std::vector<std::string>> map = get_all_open_files(); + return map[pid]; +} + +std::unordered_map<int, std::vector<std::string>> get_all_open_files() { + std::unordered_map<int, std::vector<std::string>> map; + std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> info = GetSystemHandleInformation(); + for (auto& h : info) { + int pid = h.UniqueProcessId; + + if (!IsFileHandle(nullptr, h.ObjectTypeIndex)) + continue; + if (!IsFileMaskOk(h.GrantedAccess)) + continue; + + const HANDLE proc = ::OpenProcess(PROCESS_DUP_HANDLE, false, pid); + HANDLE handle = DuplicateHandle(proc, h.HandleValue); + if (!handle) + continue; + + if (GetFileType(handle) != FILE_TYPE_DISK) + continue; + + std::string path = GetFinalPathNameByHandle(handle); + if (!IsFilePathOk(path)) + continue; + + if (map.find(pid) == map.end()) + map[pid] = {}; + map[pid].push_back(path); + } + return map; +} + +}