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

namespace pe
{
	inline PIMAGE_NT_HEADERS get_nt_headers_unsafe(uint8_t* base)
	{
		const auto dos_header = PIMAGE_DOS_HEADER(base);
		const auto nt_headers = PIMAGE_NT_HEADERS(base + dos_header->e_lfanew);

		return nt_headers;
	}

	PIMAGE_NT_HEADERS get_nt_headers(uint8_t* base)
	{
		const auto dos_header = PIMAGE_DOS_HEADER(base);
		if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)
			return nullptr;

		const auto nt_headers = PIMAGE_NT_HEADERS(base + dos_header->e_lfanew);
		if (nt_headers->Signature != IMAGE_NT_SIGNATURE || nt_headers->FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64)
			return nullptr;

		return nt_headers;
	}

	template <typename T>
	T* file_rva_to_va(uint8_t* base, const uint32_t rva)
	{
		const auto nt_headers = get_nt_headers_unsafe(base);
		const auto sections = IMAGE_FIRST_SECTION(nt_headers);

		for (size_t i = 0; i < nt_headers->FileHeader.NumberOfSections; ++i)
		{
			const auto& section = sections[i];
			if (rva >= section.VirtualAddress && rva < section.VirtualAddress + section.SizeOfRawData)
				return (T*)(base + (rva - section.VirtualAddress + section.PointerToRawData));
		}

		return (T*)(base + rva);
	}
}

std::vector<uint8_t> read_file(const std::string& file_path)
{
	std::ifstream stream(file_path, std::ios::in | std::ios::ate | std::ios::binary);
	if (!stream)
		return {};

	const auto size = stream.tellg();
	stream.seekg(0, std::ios::beg);

	std::vector<uint8_t> buffer(size);
	stream.read((char*)buffer.data(), size);

	return buffer;
}

void check_file(const std::string& file_path)
{
	auto buffer = read_file(file_path);

	if (buffer.empty())
	{
		std::printf("Failed to read file: %s.\n", file_path.c_str());
		return;
	}

	const auto base = buffer.data();
	const auto nt_headers = pe::get_nt_headers(base);
	if (!nt_headers || nt_headers->OptionalHeader.Subsystem != IMAGE_SUBSYSTEM_NATIVE)
		return;

	const auto import_dir = nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
	if (!import_dir.Size || !import_dir.VirtualAddress)
		return;

	bool printed_name = false;
	auto import_descriptor = pe::file_rva_to_va<IMAGE_IMPORT_DESCRIPTOR>(base, import_dir.VirtualAddress);
	for (; import_descriptor->Name; ++import_descriptor)
	{
		const auto module_name = pe::file_rva_to_va<const char>(base, import_descriptor->Name);
		auto imported_func = pe::file_rva_to_va<IMAGE_THUNK_DATA>(base,
			import_descriptor->OriginalFirstThunk ? import_descriptor->OriginalFirstThunk : import_descriptor->FirstThunk);

		for (; imported_func->u1.AddressOfData; ++imported_func)
		{
			if (imported_func->u1.Ordinal & IMAGE_ORDINAL_FLAG)
				continue;

			const auto import_name = pe::file_rva_to_va<IMAGE_IMPORT_BY_NAME>(base,
				uint32_t(imported_func->u1.AddressOfData))->Name;
			std::printf("%s,", import_name);
		}
	}
}

int __cdecl main(int argc, char** argv)
{
	if (argc < 2)
	{
		std::printf("[!] please provide a path to a PE file...");
		return -1;
	}
	check_file(argv[1]);
}