#include <iostream>
#include <Windows.h>
#include <fstream>
#include <filesystem>

#include <vmp2.hpp>
#include <vmprofiler.hpp>
#include <cli-parser.hpp>
#include <xtils.hpp>

int __cdecl main(int argc, const char* argv[])
{
	argparse::argument_parser_t parser(
		"vmprofiler-cli", "virtual instruction pseudo code generator");

	parser.add_argument()
		.names({ "--bin", "--vmpbin" })
		.description("unpacked binary protected with VMProtect 2")
		.required(true);

	parser.add_argument()
		.names({ "--vmentry", "--entry" })
		.description("rva to push prior to a vm_entry")
		.required(true);

	parser.add_argument()
		.name("--showhandlers")
		.description("show all vm handlers...");

	parser.enable_help();
	auto err = parser.parse(argc, argv);

	if (err)
	{
		std::cout << err << std::endl;
		return -1;
	}

	if (parser.exists("help"))
	{
		parser.print_help();
		return 0;
	}

	const auto module_base = 
		reinterpret_cast<std::uintptr_t>(
			LoadLibraryExA(parser.get<std::string>("bin").c_str(),
				NULL, DONT_RESOLVE_DLL_REFERENCES));

	const auto vm_entry_ptr = 
		module_base + std::strtoull(
			parser.get<std::string>("vmentry").c_str(), nullptr, 16);

	const auto image_base = 
		xtils::um_t::get_instance()->image_base(
			parser.get<std::string>("bin").c_str());

	zydis_routine_t vm_entry;
	std::printf("> vm entry start = 0x%p\n", vm_entry_ptr);

	if (!vm::util::flatten(vm_entry, vm_entry_ptr))
	{
		std::printf("> failed to flatten vm entry...\n");
		return -1;
	}

	vm::util::deobfuscate(vm_entry);
	std::printf("> flattened vm entry...\n");
	std::printf("> deobfuscated vm entry...\n");
	std::printf("==================================================================================\n");
	vm::util::print(vm_entry);

	const auto vm_handler_table = 
		vm::handler::table::get(vm_entry);

	if (!vm_handler_table)
	{
		std::printf("> failed to locate vm handler table...\n");
		return -1;
	}

	std::printf("==================================================================================\n");
	std::printf("> located vm handler table... at = 0x%p, rva = 0x%p\n", vm_handler_table,
		(reinterpret_cast<std::uintptr_t>(vm_handler_table) - module_base) + image_base);

	zydis_decoded_instr_t vm_handler_transform;
	if (!vm::handler::table::get_transform(vm_entry, &vm_handler_transform))
	{
		std::printf("[!] failed to locate vm handler table entry transformation...\n");
		return -1;
	}

	std::printf("> vm handler table entries decrypted with = ");
	vm::util::print(vm_handler_transform);

	vm_handler_transform.mnemonic = vm::transform::inverse[vm_handler_transform.mnemonic];
	std::printf("> vm handler table entries encrypted with = ");
	vm::util::print(vm_handler_transform);

	std::printf("==================================================================================\n");
	std::vector<zydis_decoded_instr_t> vinstr_rva_decrypt_instrs;
	if (!vm::instrs::get_rva_decrypt(vm_entry, vinstr_rva_decrypt_instrs))
	{
		std::printf("[!] failed to get virtual instruction rva decrypt transformations...\n");
		return -1;
	}

	std::printf("> virtual instruction rva decryption instructions:\n");
	for (auto& transform : vinstr_rva_decrypt_instrs)
	{
		std::printf("\t");
		vm::util::print(transform);
	}

	std::vector<vm::handler::handler_t> vm_handlers;
	if (!vm::handler::get_all(module_base, image_base, vm_entry, vm_handler_table, vm_handlers))
	{
		std::printf("> failed to get all vm handler meta data...\n");
		return -1;
	}

	if (parser.exists("showhandlers"))
	{
		for (auto idx = 0u; idx < vm_handlers.size(); ++idx)
		{
			auto vm_handler = vm_handlers[idx];
			std::printf("==========[vm handler %s, idx = %d, imm size = %d]========\n",
				vm_handler.profile ? vm_handler.profile->name : "UNKNOWN", idx,
				vm_handler.imm_size);

			std::printf("================[vm handler instructions]==============\n");
			vm::util::print(vm_handler.instrs);

			if (vm_handler.imm_size)
			{
				std::printf("=================[vm handler transforms]===============\n");
				for (const auto& [transform_type, transform] : vm_handler.transforms)
					vm::util::print(transform);
			}

			std::printf("=======================================================\n\n");
		}
	}
	std::printf("> finished...\n");
}