#include "theo.h"

namespace theo
{
	hmm_ctx::hmm_ctx(const mapper_routines_t& routines)
		:
		kalloc(std::get<0>(routines)),
		kmemcpy(std::get<1>(routines)),
		resolve_symbol(std::get<2>(routines))
	{}

	auto hmm_ctx::map_objs(std::vector<lnk::obj_buffer_t>& objs) -> image_entry_t
	{
		std::printf("[+] allocating space for symbols...\n");
		if (!alloc_symbol_space(objs))
		{
			std::printf("[!] failed to allocate symbol space...\n");
			return {};
		}

		std::printf("[+] allocating space for obfuscated symbols...\n");
		if (!alloc_obfuscated_symbol_space(objs))
		{
			std::printf("[!] failed to allocate space for obfuscated functions...\n");
			return {};
		}

		std::printf("[+] mapping obfuscated symbols...\n");
		if (!map_obfuscated_symbols(objs))
		{
			std::printf("[!] failed to resolve obfuscated relocs...\n");
			return {};
		}

		std::printf("[+] resolving non-obfuscated relocations...\n");
		if (!resolve_relocs(objs))
		{
			std::printf("[!] failed to resolve relocations...\n");
			return {};
		}

		std::printf("[+] mapping non-obfuscated symbols...\n");
		if (!map_symbols(objs))
		{
			std::printf("> failed to map symbols into memory...\n");
			return {};
		}

		std::printf("[+] linking complete...\n");
		return mapped_symbols["DrvEntry"];
	}

	bool hmm_ctx::map_symbols(std::vector<lnk::obj_buffer_t>& objs)
	{
		for (auto& obj : objs)
		{
			for (auto& symbol : lnk::sym::get_all(obj))
			{
				// dont map obfuscated routines into memory as they
				// get mapped differently...
				if (symbol.obfuscate_type)
					continue;

				const auto symbol_mapped =
					reinterpret_cast<void*>(
						mapped_symbols[symbol.symbol_name]);

				if (!symbol_mapped)
				{
					std::printf("	> failed to resolve symbol allocation = %s\n", 
						symbol.symbol_name.c_str());

					return false;
				}

				std::printf("	> mapping symbol = %s, at = 0x%p\n", 
					symbol.symbol_name.c_str(), symbol_mapped);

				kmemcpy(symbol_mapped, obj.data() + 
					symbol.file_offset, symbol.size);
			}
		}
		return true;
	}

	bool hmm_ctx::resolve_relocs(std::vector<lnk::obj_buffer_t>& objs)
	{
		for (auto& obj : objs)
		{
			for (auto& reloc : lnk::sym::get_relocs(obj))
			{
				if (reloc.type != IMAGE_REL_AMD64_ADDR64)
				{
					std::printf("[!] error... unsupported relocation at file offset = 0x%x\n\t> symbol = %s\n",
						reloc.file_offset, reloc.resolve_symbol_name.c_str());

					return false;
				}

				const auto reloc_addr =
					reinterpret_cast<std::uintptr_t*>(
						obj.data() + reloc.file_offset);

				// check obj symbol table for this relocation...
				if (mapped_symbols[reloc.resolve_symbol_name])
				{
					std::printf("	> resolving internal symbol...\n");
					std::printf("		> address = 0x%p\n", mapped_symbols[reloc.resolve_symbol_name]);
					std::printf("		> symbol = %s\n", reloc.resolve_symbol_name.c_str());
					*reloc_addr = mapped_symbols[reloc.resolve_symbol_name];
				}
				else // else check external symbol table...
				{
					const auto extern_symbol =
						resolve_symbol(reloc.resolve_symbol_name.c_str());

					if (!extern_symbol)
					{
						std::printf("[!] unresolved external symbol = %s...\n",
							reloc.resolve_symbol_name.c_str());

						return false;
					}

					*reloc_addr = extern_symbol;
				}

				std::printf("	> resolving external symbol...\n");
				std::printf("		> address = 0x%p\n", *reloc_addr);
				std::printf("		> symbol = %s\n", reloc.resolve_symbol_name.c_str());
			}
		}
		return true;
	}

	bool hmm_ctx::map_obfuscated_symbols(std::vector<lnk::obj_buffer_t>& objs)
	{
		for (auto& obj : objs)
		{
			for (auto& symbol : lnk::sym::get_all(obj))
			{
				if (!symbol.obfuscate_type)
					continue;
				
				std::printf("	> mapping obfuscated routine %s into memory...\n", symbol.symbol_name.c_str());
				std::int32_t instruc_offset = 0u;

				while (true) // TODO: this is bad code... dont do this!
				{
					auto symbol_name = symbol.symbol_name;

					if (instruc_offset)
						symbol_name.append("@")
							.append(std::to_string(instruc_offset));

					// if there is no allocation for this symbol then we are done...
					if (!mapped_symbols[symbol_name])
						break;

					const auto instruc_len = obfuscated_gadgets
						[mapped_symbols[symbol_name]]->get_instruc().length;

					auto gadget_stack = obfuscated_gadgets
						[mapped_symbols[symbol_name]]->get_gadget();

					const auto gadget_size = obfuscated_gadgets
						[mapped_symbols[symbol_name]]->get_size();

					unsigned gadget_offset = 0u;
					std::vector<std::uint8_t> gadget_raw;

					for (auto& [gadget, reloc] : gadget_stack)
					{
						const auto fix_reloc_addr = 
							gadget.data() + reloc.offset;

						switch (reloc.type)
						{
							case obfuscation::reloc_type::jcc:
							{
								const auto next_instruc_symbol =
									std::string(symbol.symbol_name).append("@")
										.append(std::to_string(instruc_offset + instruc_len + reloc.rva));

								*reinterpret_cast<std::uintptr_t*>(fix_reloc_addr) =
									mapped_symbols[next_instruc_symbol];
								break;
							}
							case obfuscation::reloc_type::next_instruction_addr:
							{
								const auto next_instruc_symbol =
									std::string(symbol.symbol_name).append("@")
										.append(std::to_string(instruc_offset + instruc_len));

								*reinterpret_cast<std::uintptr_t*>(fix_reloc_addr) = 
									mapped_symbols[next_instruc_symbol];

								break; // we resolved our own relocation...
							}
							case obfuscation::reloc_type::none:
							{
								break;
							}
							default: 
							{
								// check this instruction to see if it needs any relocs...
								for (auto& reloc : lnk::sym::get_relocs(obj))
								{
									if (reloc.file_offset >= symbol.file_offset + instruc_offset &&
										reloc.file_offset < symbol.file_offset + instruc_offset + instruc_len)
									{
										std::printf("		> resolving relocation for instruction...\n");
										if (reloc.type != IMAGE_REL_AMD64_ADDR64)
										{
											std::printf("[!] error, cannot resolve reloc = %s, type = 0x%x\n",
												reloc.resolve_symbol_name.c_str(), reloc.type);

											// cant relocate anything but IMAGE_REL_AMD64_ADDR64...
											// this is fine since the compiler shouldnt ever make any rip relative code
											// besides JCC's...
											return false;
										}

										const auto reloc_instruc_offset =
											reloc.file_offset - (symbol.file_offset + instruc_offset);

										const auto reloc_addr =
											reinterpret_cast<std::uintptr_t*>(
												gadget.data() + reloc_instruc_offset);

										// check obj symbol table for this relocation...
										if (mapped_symbols[reloc.resolve_symbol_name])
										{
											*reloc_addr = mapped_symbols[reloc.resolve_symbol_name];
										}
										else // else check external symbol table...
										{
											const auto extern_symbol =
												resolve_symbol(reloc.resolve_symbol_name.c_str());

											if (!extern_symbol)
											{
												std::printf("[!] unresolved external symbol = %s...\n",
													reloc.resolve_symbol_name.c_str());

												return false;
											}

											*reloc_addr = extern_symbol;
										}

										std::printf("			> address = 0x%p\n", *reloc_addr);
										std::printf("			> symbol = %s\n", reloc.resolve_symbol_name.c_str());
										break; // break out of for loop... we resolve the symbol...
									}
								}
							}
						}

						gadget_raw.insert(gadget_raw.end(), gadget.begin(), gadget.end());
						gadget_offset += gadget.size();
					}

					const auto gadget_addr = 
						reinterpret_cast<void*>(
							mapped_symbols[symbol_name]);

					std::printf("	> copying gadget at = 0x%p\n", gadget_addr);
					kmemcpy(gadget_addr, gadget_raw.data(), gadget_raw.size());
					// used to calc symbol for next instruction...
					instruc_offset += instruc_len;
				}
			}
		}
		return true;
	}

	bool hmm_ctx::alloc_obfuscated_symbol_space(std::vector<lnk::obj_buffer_t>& objs)
	{
		ZydisDecoder decoder;
		ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64);

		for (auto& obj : objs)
		{
			for (auto& symbol : lnk::sym::get_all(obj))
			{
				// skip normal routines for now... those get scattered...
				if (!symbol.obfuscate_type)
					continue;

				ZyanUSize offset = 0;
				ZyanUSize length = symbol.size;
				ZydisDecodedInstruction instruction;

				const auto routine_begin = 
					symbol.file_offset + obj.data();

				bool first_instruction = true;
				while (ZYAN_SUCCESS(ZydisDecoderDecodeBuffer(
					&decoder, routine_begin + offset, length - offset, &instruction)))
				{
					// dont append @offset for the first instruction...
					auto new_symbol = symbol.symbol_name;

					if (first_instruction)
						first_instruction = false;
					else
						new_symbol.append("@")
							.append(std::to_string(offset));

					std::vector<std::uint8_t> instruc_bytes{};
					instruc_bytes.resize(instruction.length);

					memcpy(instruc_bytes.data(), obj.data() +
						symbol.file_offset + offset, instruction.length);

					std::shared_ptr<obfuscation::obfuscate> new_gadget{};
					switch (symbol.obfuscate_type)
					{
						case lnk::theo_type::obfuscate:
						{
							new_gadget.reset(
								new obfuscation::obfuscate(
									{ instruction, instruc_bytes }));
							break;
						}
						case lnk::theo_type::mutate:
						{
							new_gadget.reset(
								new obfuscation::mutation(
									{ instruction, instruc_bytes }));
							break;
						}
						default:
						{
							std::printf("[!] unsupported obfuscation type on routine = %s, type = %d\n", 
								symbol.symbol_name.c_str(), symbol.obfuscate_type);
							return false;
						}
					}

					mapped_symbols[new_symbol] =
						reinterpret_cast<std::uintptr_t>(
							kalloc(new_gadget->get_size()));

					obfuscated_gadgets[mapped_symbols[new_symbol]] = new_gadget;
					std::printf("		> %s allocated = 0x%p, size = %d\n", new_symbol.c_str(),
						mapped_symbols[new_symbol], new_gadget->get_size());

					offset += instruction.length;
				}
			}
		}
		return true;
	}

	bool hmm_ctx::alloc_symbol_space(std::vector<lnk::obj_buffer_t>& objs)
	{
		for (auto& obj : objs)
		{
			for (auto& symbol : lnk::sym::get_all(obj))
			{
				// skip obfuscated routines for now... those get scattered...
				if (symbol.obfuscate_type)
					continue;

				mapped_symbols[symbol.symbol_name] =
					reinterpret_cast<std::uintptr_t>(kalloc(symbol.size));

				std::printf("	> %s allocated at = 0x%p, size = %d\n",
					symbol.symbol_name.c_str(), mapped_symbols[symbol.symbol_name], symbol.size);
			}
		}
		return true;
	}
}