#include "pch.h"
#include "Processes.h"
#include <Psapi.h>
#include <VersionHelpers.h>
#include <array>
#include <assert.h>
#include "Thread.h"

using namespace WinSys;

namespace WinSys {
	enum class ProcessProtectionSigner :uint8_t {
		None = PsProtectedSignerNone,
		Authenticode = PsProtectedSignerAuthenticode,
		CodeGen = PsProtectedSignerCodeGen,
		Antimalware = PsProtectedSignerAntimalware,
		Lsa = PsProtectedSignerLsa,
		Windows = PsProtectedSignerWindows,
		WinTcb = PsProtectedSignerWinTcb,
		WinSystem = PsProtectedSignerWinSystem,
		App = PsProtectedSignerApp,
		Max = PsProtectedSignerMax
	};
}

static bool GetProcessPeb(HANDLE hProcess, PPEB peb) {
	ULONG len = 0;
	PROCESS_BASIC_INFORMATION info;
	DWORD status = ::NtQueryInformationProcess(hProcess, ProcessBasicInformation, &info, sizeof(info), &len);
	if(!NT_SUCCESS(status))
		return false;

	if (info.PebBaseAddress == nullptr)
		return false;

	return ::ReadProcessMemory(hProcess, info.PebBaseAddress, peb, sizeof(*peb), nullptr);
}

static bool GetExtendedInfo(HANDLE hProcess, PROCESS_EXTENDED_BASIC_INFORMATION* info) {
	ULONG len;
	info->Size = sizeof(*info);
	auto status = ::NtQueryInformationProcess(hProcess, ProcessBasicInformation, info, sizeof(*info), &len);
	return NT_SUCCESS(status);
}

uint32_t Process::GetId() const {
	return ::GetProcessId(_handle.get());
}

std::wstring Process::GetName() const {
	auto name = GetFullImageName();
	auto slash = name.rfind(L'\\');
	return slash == std::wstring::npos ? name : name.substr(slash + 1);
}

std::optional<ProcessWindowInfo> Process::GetWindowInformation() const {
	auto buffer = std::make_unique<BYTE[]>(1024);
	std::optional<ProcessWindowInfo> info;
	ULONG len;
	if (!NT_SUCCESS(::NtQueryInformationProcess(_handle.get(), ProcessWindowInformation, buffer.get(), 1024, &len)))
		return info;

	auto p = reinterpret_cast<PROCESS_WINDOW_INFORMATION*>(buffer.get());
	info->Flags = p->WindowFlags;
	info->Title = std::wstring(p->WindowTitle, p->WindowTitleLength);
	return info;
}

int Process::GetMemoryPriority() const {
	int priority = -1;
	ULONG len;
	::NtQueryInformationProcess(_handle.get(), ProcessPagePriority, &priority, sizeof(priority), &len);
	return priority;
}

IoPriority Process::GetIoPriority() const {
	IoPriority priority = IoPriority::Unknown;
	ULONG len;
	::NtQueryInformationProcess(_handle.get(), ProcessIoPriority, &priority, sizeof(priority), &len);
	return priority;
}

HANDLE Process::GetHandle() const {
	return _handle.get();
}

bool Process::IsElevated() const {
	wil::unique_handle hToken;
	if (!::OpenProcessToken(_handle.get(), TOKEN_QUERY, hToken.addressof()))
		return false;

	TOKEN_ELEVATION elevation;
	DWORD size;
	if (!::GetTokenInformation(hToken.get(), TokenElevation, &elevation, sizeof(elevation), &size))
		return false;

	return elevation.TokenIsElevated ? true : false;
}

std::wstring Process::GetWindowTitle() const {
	BYTE buffer[1024];
	ULONG len;
	auto status = ::NtQueryInformationProcess(_handle.get(), ProcessWindowInformation, buffer, 1024, &len);
	if (!NT_SUCCESS(status))
		return L"";

	auto name = reinterpret_cast<PROCESS_WINDOW_INFORMATION*>(buffer);
	return std::wstring(name->WindowTitle, name->WindowTitleLength / sizeof(WCHAR));
}

IntegrityLevel Process::GetIntegrityLevel() const {
	wil::unique_handle hToken;
	if (!::OpenProcessToken(_handle.get(), TOKEN_QUERY, hToken.addressof()))
		return IntegrityLevel::Error;

	BYTE buffer[256];
	DWORD len;
	if (!::GetTokenInformation(hToken.get(), TokenIntegrityLevel, buffer, 256, &len))
		return IntegrityLevel::Error;

	auto integrity = reinterpret_cast<TOKEN_MANDATORY_LABEL*>(buffer);

	auto sid = integrity->Label.Sid;
	return (IntegrityLevel)(*::GetSidSubAuthority(sid, *::GetSidSubAuthorityCount(sid) - 1));
}

std::unique_ptr<Process> WinSys::Process::GetCurrent() {
	return std::make_unique<Process>(NtCurrentProcess());
}

Process::Process(HANDLE handle) :_handle(handle) {
}

std::unique_ptr<Process> Process::OpenById(uint32_t pid, ProcessAccessMask access) {
	auto handle = ::OpenProcess(static_cast<ACCESS_MASK>(access), FALSE, pid);
	return handle ? std::make_unique<Process>(handle) : nullptr;
}

Process::~Process() = default;// 将该函数声明为显示默认构造函数,比手动编写函数更加有效

bool WinSys::Process::IsValid() const {
	return _handle != nullptr;
}

std::wstring Process::GetFullImageName() const {
	DWORD size = MAX_PATH;
	WCHAR name[MAX_PATH];
	auto success = ::QueryFullProcessImageName(_handle.get(), 0, name, &size);
	return success ? std::wstring(name) : L"";
}

std::wstring Process::GetCmdLine() const {
	if (::IsWindows8OrGreater()) {
		ULONG size = 8192;
		auto buffer = std::make_unique<BYTE[]>(size);
		auto status = ::NtQueryInformationProcess(_handle.get(), ProcessCommandLineInformation, buffer.get(), size, &size);
		if (NT_SUCCESS(status)) {
			auto str = (UNICODE_STRING*)buffer.get();
			return std::wstring(str->Buffer, str->Length / sizeof(WCHAR));
		}
		if (status == STATUS_INVALID_INFO_CLASS) {
			return L"";
		}
	}
	else {
		return GetCmdLine(_handle.get());
	}
	return L"";
}

std::wstring Process::GetTokenUserName() const {
	wil::unique_handle hToken;
	if (!::OpenProcessToken(_handle.get(), TOKEN_QUERY, hToken.addressof()))
		return L"";

	BYTE buffer[128];
	DWORD len;
	if (!::GetTokenInformation(hToken.get(), TokenUser, buffer, sizeof(buffer), &len))
		return L"";

	auto user = reinterpret_cast<TOKEN_USER*>(buffer);
	DWORD userMax = TOKEN_USER_MAX_SIZE;
	wchar_t name[TOKEN_USER_MAX_SIZE];
	DWORD domainMax = 64;
	wchar_t domain[64];
	SID_NAME_USE use;
	if (!::LookupAccountSid(nullptr, user->User.Sid, name, &userMax, domain, &domainMax, &use))
		return L"";

	return std::wstring(domain) + L"\\" + name;
}

std::optional<ProcessProtection> Process::GetProtection() const {
	ProcessProtection protection;
	ULONG len;
	auto status = ::NtQueryInformationProcess(_handle.get(), ProcessProtectionInformation, &protection, sizeof(protection), &len);
	if (!NT_SUCCESS(status))
		return std::optional<ProcessProtection>();
	return protection;
}

bool Process::Terminate(uint32_t exitCode) {
	return NT_SUCCESS(::NtTerminateProcess(_handle.get(), exitCode));
}

bool Process::Suspend() {
	return NT_SUCCESS(::NtSuspendProcess(_handle.get()));
}

bool Process::Resume() {
	return NT_SUCCESS(::NtResumeProcess(_handle.get()));
}

bool Process::IsImmersive() const noexcept {
	using PIsImmersiveProcess = BOOL (WINAPI*) (HANDLE);
	// Determines whether the process belongs to a Windows Store app.
	if (IsWindows8OrGreater()) {
		PIsImmersiveProcess pIsImmersiveProcess = (PIsImmersiveProcess)GetProcAddress(GetModuleHandle(L"user32.dll"), "IsImmersiveProcess");
		return pIsImmersiveProcess(_handle.get()) ? true : false;
	}
	return false;
}

bool Process::IsProtected() const {
	PROCESS_EXTENDED_BASIC_INFORMATION info;
	if (!GetExtendedInfo(GetHandle(), &info))
		return false;
	
	return info.IsProtectedProcess ? true : false;
}

bool Process::IsSecure() const {
	PROCESS_EXTENDED_BASIC_INFORMATION info;
	if(!GetExtendedInfo(GetHandle(),&info))
		return false;

	return info.IsSecureProcess ? true : false;
}

bool Process::IsInJob(HANDLE hJob) const {
	BOOL injob = FALSE;
	::IsProcessInJob(_handle.get(), hJob, &injob);
	return injob ? true : false;
}

bool Process::IsWow64Process() const {
	PROCESS_EXTENDED_BASIC_INFORMATION info;
	if (!GetExtendedInfo(GetHandle(), &info))
		return false;

	return info.IsWow64Process ? true : false;
}

bool Process::IsManaged() const {
	wil::unique_handle hProcess;
	if (!::DuplicateHandle(::GetCurrentProcess(), _handle.get(), ::GetCurrentProcess(), hProcess.addressof(),
		PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, 0))
		return false;

	WCHAR filename[MAX_PATH], sysPath[MAX_PATH];
	BOOL wow64 = FALSE;
	::IsWow64Process(hProcess.get(), &wow64);
	::GetSystemDirectory(sysPath, MAX_PATH);
	::wcscat_s(sysPath, L"\\mscoree.dll");

	HMODULE hModule[64];
	DWORD needed;
	if (!::EnumProcessModulesEx(hProcess.get(), hModule, sizeof(hModule), &needed, wow64 ? LIST_MODULES_32BIT : LIST_MODULES_ALL))
		return false;

	int count = min(_countof(hModule), needed / sizeof(HMODULE));

	for (int i = 0; i < count; i++) {
		if(::GetModuleFileNameEx(hProcess.get(),hModule[i],filename,MAX_PATH)==0)
			continue;
		if (::_wcsicmp(filename, sysPath) == 0)
			return true;// 是否是托管程序
	}
	return false;
}

WinSys::ProcessPriorityClass Process::GetPriorityClass() const {
	return static_cast<ProcessPriorityClass>(::GetPriorityClass(_handle.get()));
}

bool WinSys::Process::SetPriorityClass(ProcessPriorityClass pc) {
	return ::SetPriorityClass(_handle.get(), static_cast<DWORD>(pc));
}

uint32_t WinSys::Process::GetGdiObjectCount() const {
	return ::GetGuiResources(_handle.get(), GR_GDIOBJECTS);
}

uint32_t WinSys::Process::GetPeakGdiObjectCount() const {
	return ::GetGuiResources(_handle.get(), GR_GDIOBJECTS_PEAK);
}

uint32_t WinSys::Process::GetUserObjectCount() const {
	return ::GetGuiResources(_handle.get(), GR_USEROBJECTS);
}

uint32_t WinSys::Process::GetPeakUserObjectCount() const {
	return ::GetGuiResources(_handle.get(), GR_USEROBJECTS_PEAK);
}

HANDLE WinSys::Process::GetNextThread(HANDLE hThread, ThreadAccessMask access) {
	HANDLE hNewThread{ nullptr };
	::NtGetNextThread(_handle.get(), hThread, static_cast<ACCESS_MASK>(access), 0, 0, &hNewThread);
	return hNewThread;
}

std::wstring WinSys::Process::GetCurDirectory() const {
	wil::unique_handle h;
	if (!::DuplicateHandle(::GetCurrentProcess(), _handle.get(), ::GetCurrentProcess(), h.addressof(),
		PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, 0))
		return L"";
	return GetCurDirectory(h.get());
}

std::wstring Process::GetCurDirectory(HANDLE hProcess) {
	std::wstring path;
	PEB peb;
	if (!GetProcessPeb(hProcess, &peb))
		return path;

	RTL_USER_PROCESS_PARAMETERS processParams;
	if (!::ReadProcessMemory(hProcess, peb.ProcessParameters, &processParams, sizeof(processParams), nullptr))
		return path;

	path.resize(processParams.CurrentDirectory.DosPath.Length / sizeof(WCHAR) + 1);
	if (!::ReadProcessMemory(hProcess, processParams.CurrentDirectory.DosPath.Buffer, path.data(),
		processParams.CurrentDirectory.DosPath.Length, nullptr))
		return L"";

	return path;
}

std::wstring Process::GetCmdLine(HANDLE hProcess) {
	std::wstring cmdLine;
	PEB peb;
	if (!GetProcessPeb(hProcess, &peb))
		return cmdLine;

	RTL_USER_PROCESS_PARAMETERS processParams;
	if (!::ReadProcessMemory(hProcess, peb.ProcessParameters, &processParams, sizeof(processParams), nullptr))
		return cmdLine;

	cmdLine.resize(processParams.CommandLine.Length / sizeof(WCHAR) + 1);
	if (!::ReadProcessMemory(hProcess, processParams.CommandLine.Buffer, cmdLine.data(),
		processParams.CommandLine.Length, nullptr))
		return L"";

	return cmdLine;
}

std::vector<std::pair<std::wstring, std::wstring>> Process::GetEnvironment(HANDLE hProcess) {
	std::vector<std::pair<std::wstring, std::wstring>> env;

	PEB peb;
	if (!GetProcessPeb(hProcess, &peb))
		return env;

	RTL_USER_PROCESS_PARAMETERS processParams;
	if (!::ReadProcessMemory(hProcess, peb.ProcessParameters, &processParams, sizeof(processParams), nullptr))
		return env;

	BYTE* buffer = nullptr;
	int size = 1 << 16;
	buffer = (BYTE*)malloc(size);
	if (!buffer) {
		return env;
	}
	for (;;) {
		if (::ReadProcessMemory(hProcess, processParams.Environment, buffer, size, nullptr))
			break;
		size -= 1 << 12;
	}

	env.reserve(64);

	for (auto p = (PWSTR)buffer; *p;) {
		std::pair<std::wstring, std::wstring> var;
		auto equal = wcschr(p, L'=');
		assert(equal);
		if (!equal)
			break;
		*equal = L'\0';
		var.first = p;
		p += ::wcslen(p) + 1;
		var.second = p;
		p += ::wcslen(p) + 1;
		env.push_back(std::move(var));
	}

	free(buffer);

	return env;
}


DWORD_PTR Process::GetImageBaseAddress(HANDLE hProcess) {
	PEB peb;
	if (!GetProcessPeb(hProcess, &peb))
		return 0;

	return (DWORD_PTR)peb.ImageBaseAddress;
}

SIZE_T Process::GetImageSize(HANDLE hProcess, DWORD_PTR imageBase) {
	MEMORY_REGION_INFORMATION info = { 0 };
	SIZE_T size = 0;
	SIZE_T len = 0;
	DWORD status = ::NtQueryVirtualMemory(hProcess, (PVOID)imageBase, MemoryRegionInformation,
		&info, sizeof(MEMORY_REGION_INFORMATION), &len);
	if (NT_SUCCESS(status)) {
		return info.RegionSize;
	}
	return 0;
}

bool Process::HasVEH(HANDLE hProcess) const {
	ULONG len = 0;
	PROCESS_BASIC_INFORMATION info;
	DWORD status = ::NtQueryInformationProcess(hProcess, ProcessBasicInformation, &info, sizeof(info), &len);
	if (!NT_SUCCESS(status))
		return false;

	if (info.PebBaseAddress == nullptr)
		return false;
	DWORD flags = 0;
	DWORD* pCrossProcessFlags = nullptr;
	PROCESS_EXTENDED_BASIC_INFORMATION extInfo;
	if (!GetExtendedInfo(hProcess, &extInfo))
		return false;
	if (extInfo.IsWow64Process) {
		pCrossProcessFlags = reinterpret_cast<DWORD*>((BYTE*)info.PebBaseAddress + 0x1000 + 0x28);
	}
	else {
		pCrossProcessFlags = reinterpret_cast<DWORD*>((BYTE*)info.PebBaseAddress + 0x50);
	}

	if (!::ReadProcessMemory(hProcess, pCrossProcessFlags, &flags, sizeof(flags), nullptr))
		return false;

	if (flags & 0x00000004)
		return true;

	return false;
}