#include "pch.h"
#include "Utils.h"
#include "NovaHypervisor.h"

int main(int argc, char* argv[]) {
	if (argc < 3 || argc > 5) {
		std::cerr << "Usage: NovaClient.exe <protect | unprotect> <address> <permissions [r|w|x]>" << std::endl;
		return EXIT_FAILURE;
	}
	std::string operation = argv[1];
	std::string rawAddress(argv[2]);

	if (operation.compare("protect") != 0 && operation.compare("unprotect") != 0) {
		std::cerr << "Invalid operation" << std::endl;
		return EXIT_FAILURE;
	}

	try {
		NovaHypervisor novaHypervisor = NovaHypervisor();
		UINT64 address = GetAddress(rawAddress);

		if (operation.compare("protect") == 0) {
			ProtectedMemory protectedMemory = { 0 };

			if (argc < 4) {
				std::cerr << "Missing permissions argument" << std::endl;
				return EXIT_FAILURE;
			}
			protectedMemory.Address = address;
			protectedMemory.Permissions = novaHypervisor.TranslatePermissions(argv[3]);
			
			if (protectedMemory.Permissions & EPT_PAGE_WRITE && !(protectedMemory.Permissions & EPT_PAGE_READ)) {
				std::cerr << "Cannot set write bit and not read bit" << std::endl;
				return EXIT_FAILURE;
			}

			if (novaHypervisor.ProtectAddressRange(protectedMemory)) {
				std::cout << "Protected address: 0x" << std::hex << protectedMemory.Address << " with permissions: " 
					<< argv[3] << std::endl;
				return EXIT_SUCCESS;
			}
			std::cerr << "Failed to protect address" << std::endl;
			return EXIT_FAILURE;
		}
		else if (operation.compare("unprotect") == 0) {
			if (novaHypervisor.UnprotectAddressRange(address)) {
				std::cout << "Unprotected address: 0x" << std::hex << address << std::endl;
				return EXIT_SUCCESS;
			}
			std::cerr << "Failed to protect address" << std::endl;
			return EXIT_FAILURE;
		}
	}
	catch (const std::runtime_error& e) {
		std::cerr << e.what() << std::endl;
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}