#ifndef DRIVER_CPP
#define DRIVER_CPP

#include "driver.hpp"
#include "../utils/Authentication/skStr.h"

auto driver::communicate_t::get_process_pid(
	const std::wstring& proc_name) -> const std::uint32_t
{
	PROCESSENTRY32 proc_info;
	proc_info.dwSize = sizeof(proc_info);

	HANDLE proc_snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL);
	if (proc_snapshot == INVALID_HANDLE_VALUE) {
		return 0;
	}

	Process32First(proc_snapshot, &proc_info);
	if (!wcscmp(proc_info.szExeFile, proc_name.c_str()))
	{
		CloseHandle(proc_snapshot);
		return proc_info.th32ProcessID;
	}

	while (Process32Next(proc_snapshot, &proc_info))
	{
		if (!wcscmp(proc_info.szExeFile, proc_name.c_str()))
		{
			CloseHandle(proc_snapshot);
			return proc_info.th32ProcessID;
		}
	}

	CloseHandle(proc_snapshot);
	return 0;
}

auto driver::communicate_t::send_cmd(
	void* data,
	requests code) -> bool
{
	if (!data || !code)
		return false;

	IO_STATUS_BLOCK block;
	invoke_data request{ 0 };

	request.unique = requests::invoke_unique;
	request.data = data;
	request.code = code;

	auto result =
		direct_device_control(
			this->m_handle,
			nullptr,
			nullptr,
			nullptr,
			&block,
			0,
			&request,
			sizeof(request),
			&request,
			sizeof(request));

	return result == 0;
}

auto driver::communicate_t::initialize_handle() -> bool
{
	this->m_handle = CreateFileA(device_name, GENERIC_READ, 0, 0, 3, 0x00000080, 0);
	if (this->m_handle != INVALID_HANDLE_VALUE) {
		return true;
	}
	return false;
}

auto driver::communicate_t::attach(
	int a_pid) -> bool
{
	if (!a_pid)
		return false;

	this->m_pid = a_pid;

	this->dtb = this->get_dtb(a_pid);
	if (!this->dtb) {
		std::cout << skCrypt("[!] Failed to resolve DTB for PID: ").decrypt() << a_pid << std::endl;
		return false;
	}

	std::cout << skCrypt("[+] Attached to PID: ").decrypt() << a_pid << skCrypt(" | DTB: 0x").decrypt() << std::hex << this->dtb << std::dec << std::endl;
	return true;
}


auto driver::communicate_t::get_image_base(
	const char* module_name) -> const std::uintptr_t
{
	base_invoke data{ 0 };

	data.pid = this->m_pid;
	data.handle = 0;
	data.name = module_name;

	this->send_cmd(
		&data,
		invoke_base);

	return data.handle;
}

auto driver::communicate_t::write(
	const std::uintptr_t address,
	void* buffer,
	const std::size_t size) -> bool
{
	if (!this->dtb) {
		std::cout << skCrypt("[!] DTB not resolved, cannot write!").decrypt() << std::endl;
		return false;
	}

	write_invoke data{ 0 };

	data.pid = this->m_pid;
	data.directory_table = this->dtb;
	data.address = address;
	data.buffer = buffer;
	data.size = size;

	auto result =
		this->send_cmd(
			&data,
			invoke_write);

	if (!result) {
		std::cout << skCrypt("[!] Write failed - PID: ").decrypt()
			<< this->m_pid << skCrypt(" | Addr: 0x").decrypt() << std::hex << address
			<< skCrypt(" | Size: ").decrypt() << std::dec << size << std::endl;
	}

	return result;
}

auto driver::communicate_t::read(
	const std::uintptr_t address,
	void* buffer,
	const std::size_t size) -> bool
{
	if (!this->dtb) {
		std::cout << skCrypt("[!] DTB not resolved, cannot read!").decrypt() << std::endl;
		return false;
	}

	read_invoke data{ 0 };

	data.pid = this->m_pid;
	data.dtb = this->dtb;
	data.address = address;
	data.buffer = buffer;
	data.size = size;

	auto result =
		this->send_cmd(
			&data,
			invoke_read);

	if (!result) {
		std::cout << skCrypt("[!] Read failed - PID: ").decrypt()
			<< this->m_pid << skCrypt(" | Addr: 0x").decrypt() << std::hex << address
			<< skCrypt(" | Size: ").decrypt() << std::dec << size << std::endl;
	}

	return result;
}

auto driver::communicate_t::read_kernel(
	const uintptr_t address,
	void* buffer,
	const size_t size,
	const uint32_t memory_type) -> bool
{
	read_kernel_invoke data{ 0 };

	data.address = address;
	data.buffer = buffer;
	data.size = size;
	data.memory_type = memory_type;

	auto result =
		this->send_cmd(
			&data,
			invoke_read_kernel
		);

	return result;
}

auto driver::communicate_t::translate_address(
	std::uintptr_t virtual_address,
	std::uintptr_t directory_base) -> const std::uintptr_t
{
	translate_invoke data{ 0 };

	data.virtual_address = virtual_address;
	data.directory_base = directory_base;
	data.physical_address = nullptr;

	this->send_cmd(
		&data,
		invoke_translate
	);

	return reinterpret_cast<std::uintptr_t>(data.physical_address);
}

auto driver::communicate_t::get_dtb(
	std::uint32_t pid) -> const std::uintptr_t
{
	dtb_invoke data{ 0 };

	data.pid = pid;
	data.dtb = 0;

	auto result =
		this->send_cmd(
			&data,
			invoke_dtb
		);

	if (!result) {
		std::cout << skCrypt("[!] Failed to get DTB for PID: ").decrypt() << pid << std::endl;
		return 0;
	}

	return data.dtb;
}

auto driver::communicate_t::hide_process(
	uint32_t pid) -> bool
{
	hide_process_invoke data{ 0 };
	data.process_id = pid;

	auto result = this->send_cmd(
		&data,
		invoke_hide_process
	);

	if (result) {
		std::cout << skCrypt("[+] Process hidden successfully: PID ").decrypt() << pid << std::endl;
	}
	else {
		std::cout << skCrypt("[!] Failed to hide process: PID ").decrypt() << pid << std::endl;
	}

	return result;
}

#endif