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;
+}
+
+}