/**
 * win32.cpp
 *  - provides support for Windows clients
 *
 **/
#include "win32.h"
#include <fileapi.h>
#include <handleapi.h>
#include <iostream>
#include <libloaderapi.h>
#include <ntdef.h>
#include <psapi.h>
#include <shlobj.h>
#include <stdexcept>
#include <string>
#include <stringapiset.h>
#include <tlhelp32.h>
#include <unordered_map>
#include <vector>
#include <windows.h>
#include <winternl.h>

/* 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. In fact, this thing requires
   you to use *internal functions* that can't even be linked to, hence why we have to
   use GetProcAddress and such. What a mess. */

#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 { namespace Windows {

/* 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;

		res.reserve(cb);

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

/* we're using UTF-8. originally, I had used just the ANSI versions of functions, but that
   sucks massive dick. this way we get unicode in the way every single other OS does it */
std::string UnicodeStringToUtf8(std::wstring string) {
	unsigned long size = ::WideCharToMultiByte(CP_UTF8, 0, string.c_str(), string.length(), NULL, 0, NULL, NULL);
	std::string ret = std::string(size, '\0');
	::WideCharToMultiByte(CP_UTF8, 0, string.c_str(), string.length(), &ret.front(), ret.length(), NULL, NULL);
	return ret;
}

std::string UnicodeStringToUtf8(UNICODE_STRING string) {
	unsigned long size = ::WideCharToMultiByte(CP_UTF8, 0, string.Buffer, string.Length, NULL, 0, NULL, NULL);
	std::string ret = std::string(size, '\0');
	::WideCharToMultiByte(CP_UTF8, 0, string.Buffer, string.Length, &ret.front(), ret.length(), NULL, NULL);
	return ret;
}

std::wstring Utf8StringToUnicode(std::string string) {
	unsigned long size = ::MultiByteToWideChar(CP_UTF8, 0, string.c_str(), string.length(), NULL, 0);
	std::wstring ret = std::wstring(size, L'\0');
	::MultiByteToWideChar(CP_UTF8, 0, string.c_str(), string.length(), &ret.front(), ret.length());
	return ret;
}

std::string GetHandleType(HANDLE handle) {
	OBJECT_TYPE_INFORMATION info = QueryObjectTypeInfo(handle);
	return UnicodeStringToUtf8(info.TypeName);
}

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);
	buffer.resize(buffer.find('\0'));

	return UnicodeStringToUtf8(buffer);
}

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) {
	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())
		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::string GetSystemDirectory() {
	PWSTR path_wch;
	SHGetKnownFolderPath(FOLDERID_Windows, 0, NULL, &path_wch);
	std::wstring path_wstr(path_wch);
	CoTaskMemFree(path_wch);
	return UnicodeStringToUtf8(path_wstr);
}

bool IsSystemFile(const std::string& path) {
	std::wstring path_w = Utf8StringToUnicode(path);
	CharUpperBuffW(&path_w.front(), path_w.length());
	std::wstring windir_w = Utf8StringToUnicode(GetSystemDirectory());
	CharUpperBuffW(&windir_w.front(), windir_w.length());
	return path_w.find(windir_w) == 4;
}

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) {
	unsigned long size = 256, ret_size = 0;
	HANDLE handle = ::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
	if (!handle)
		return "";

	std::wstring ret(size, L'\0');
	while (size < 32768) {
		ret.resize(size, L'\0');

		if (!(ret_size = ::GetModuleBaseNameW(handle, 0, &ret.front(), ret.length())))
			ret = L"";
		else if (size > ret_size)
			ret.resize(ret.find('\0'));

		size *= 2;
	}

	CloseHandle(handle);

	return UnicodeStringToUtf8(ret);
}

std::vector<std::string> get_open_files(int pid) {
	std::vector<std::string> ret;
	std::vector<SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX> info = GetSystemHandleInformation();
	for (auto& h : info) {
		if (h.UniqueProcessId != pid)
			continue;

		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;

		ret.push_back(path);
	}
	return ret;
}

std::vector<std::string> filter_system_files(const std::vector<std::string>& source) {
	std::vector<std::string> ret;
	for (const std::string& s : source) {
		if (!IsSystemFile(s))
			ret.push_back(s);
	}
	return ret;
}

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

} // namespace Windows
} // namespace Animia