#pragma once
#include <Windows.h>
#include <winternl.h>
#include <vector>
#include <functional>
#include <fstream>
#include "utils.hpp"

namespace lnk
{
	enum theo_type
	{
		obfuscate = 1,
		mutate = 2,
		encrypt = 3
	};

	struct symbol_t
	{
		// name of the symbol... not mangled...
		std::string symbol_name;

		// https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#type-representation
		std::uint32_t type;

		// what section this symbol is in...
		std::uint32_t section_number;

		// offset into section...
		std::uint32_t section_offset;

		// file offset into OBJ file...
		std::uint32_t file_offset;

		// only used by functions... size in bytes of routine...
		std::uint32_t size;

		// if this symbol is a function and is inside of a .theo section...
		theo_type obfuscate_type;
	};

	// redef of IMAGE_RELOCATION so that "VirtualAddress"
	// will actually be a file offset instead of a section offset...
	struct image_reloc_t
	{
		// name of the symbol to be resolved for example "ExAllocatePool"...
		std::string resolve_symbol_name;

		// offset into the obj file where the resolving needs to be done...
		std::uint32_t file_offset;

		// type of data that needs to be resolved... this will be 64bit addresses...
		// https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#type-indicators
		std::uint16_t type;
	};

	using obj_buffer_t = std::vector<std::uint8_t>;
	using map_symbols_t = std::map<std::string, std::pair<std::uint32_t, std::uint32_t>>;

	auto get_symbol_size(symbol_t& sym, obj_buffer_t& obj) -> std::uint32_t;
	auto get_objs(std::string lib_path, std::vector<obj_buffer_t>& objs) -> bool;
	auto get_map_symbols(std::string map_path) -> map_symbols_t;

	namespace sym
	{
		auto get_all(obj_buffer_t& obj) -> std::vector<symbol_t>;
		auto get_relocs(obj_buffer_t& obj)->std::vector<image_reloc_t>;
	}

	namespace section
	{
		using section_callback_t = std::function<bool(PIMAGE_SECTION_HEADER, obj_buffer_t& obj)>;
		auto get_header(obj_buffer_t& obj, const char* section_name) -> PIMAGE_SECTION_HEADER;
		auto for_each(section_callback_t callback, obj_buffer_t& obj) -> void;
	}
}