#include <iostream>
#include <Windows.h>
#include <fstream>
#include <filesystem>
#include <vmp2.hpp>
#include <vmctx.h>
#include <vm.h>
#include <cli-parser.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()
		.names({ "--imagebase", "--base" })
		.description("image base from OptionalHeader::ImageBase")
		.required(true);

	parser.add_argument()
		.names({ "--vmtrace", "--trace" })
		.description("a vmp2 file generated by a vmtracer");

	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 = std::strtoull(
		parser.get<std::string>("imagebase").c_str(), nullptr, 16);

	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");
	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("> 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);

	std::vector<vm::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;
	}

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

	if (parser.exists("vmtrace"))
	{
		std::ifstream vmp2_file(
			parser.get<std::string>("vmtrace").c_str(), std::ios::binary);

		const auto file_size =
			std::filesystem::file_size(
				parser.get<std::string>("vmtrace").c_str());

		const auto file_header =
			reinterpret_cast<vmp2::file_header*>(malloc(file_size));

		vmp2_file.read((char*)file_header, file_size);

		const auto entry_list =
			reinterpret_cast<vmp2::entry_t*>(
				reinterpret_cast<std::uintptr_t>(
					file_header) + file_header->entry_offset);

		std::printf("> module base = 0x%p\n", file_header->module_base);
		std::getchar();

		vm::vmctx_t vmctx(
			file_header,
			entry_list,
			vm_handlers,
			module_base,
			image_base
		);

		for (auto [virt_instr, log_entry] = vmctx.step(); log_entry && !virt_instr.empty();
			std::tie(virt_instr, log_entry) = vmctx.step())
				std::printf("> %s\n", virt_instr.c_str());
	}

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