#pragma once
#ifndef FUNCTIONS_H
#define FUNCTIONS_H
#define WIN32_LEAN_AND_MEAN

#include <windows.h>
#include <TlHelp32.h>
#include <vector>

struct Call
{
	int returnValue;
	std::string returnString;
	DWORD exitCode;
};

namespace functions
{
	enum class Type
	{
		T_VOID = 0x0,
		T_STRING = 0x1,
		T_CHAR = 0x2,
		T_BOOL = 0x3,
		T_INT = 0x4,
		T_DOUBLE = 0x5,
		T_FLOAT = 0x6
	};

	struct Arg
	{
		Type type;
		LPVOID value;
	};

	LPVOID reserveString(HANDLE hProcess, const char *value, SIZE_T size);
	char readChar(HANDLE hProcess, DWORD64 address);

	template <class returnDataType>
	Call call(HANDLE pHandle, std::vector<Arg> args, Type returnType, DWORD64 address, char **errorMessage)
	{
		std::vector<unsigned char> argShellcode;

		std::reverse(args.begin(), args.end());

		for (auto &arg : args)
		{
			// 0x68: PUSH imm16/imm32
			// 0x6A: PUSH imm8

			if (arg.type == Type::T_INT || arg.type == Type::T_FLOAT)
			{
				argShellcode.push_back(0x68);
				int value = *static_cast<int *>(arg.value);

				// Little endian representation
				for (int i = 0; i < 4; i++)
				{
					int shifted = (value >> (i * 8)) & 0xFF;
					argShellcode.push_back(shifted);
				}

				continue;
			}

			if (arg.type == Type::T_STRING)
			{
				argShellcode.push_back(0x68);
				std::string value = *static_cast<std::string *>(arg.value);
				LPVOID address = functions::reserveString(pHandle, value.c_str(), value.length());

				// Little endian representation
				for (int i = 0; i < 4; i++)
				{
					int shifted = ((int)address >> (i * 8)) & 0xFF;
					argShellcode.push_back(shifted);
				}

				continue;
			}

			argShellcode.push_back(0x6A);
			unsigned char value = *static_cast<unsigned char *>(arg.value);
			argShellcode.push_back(value);
		}

		// 83: ADD r/m16/32 imm8
		std::vector<unsigned char> callShellcode = {
			// call 0x00000000
			0xE8,
			0x00,
			0x00,
			0x00,
			0x00,
			// add esp, [arg count * 4]
			0x83,
			0xC4,
			(unsigned char)(args.size() * 0x4),
		};

		LPVOID returnValuePointer = 0;
		if (returnType != Type::T_VOID)
		{
			// We will reserve memory for where we want to store the result,
			// and move the return value to this address.
			returnValuePointer = VirtualAllocEx(pHandle, NULL, sizeof(returnDataType), MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);

			if (returnType == Type::T_FLOAT)
			{
				// fstp DWORD PTR [0x12345678]
				// when `call` is executed, if return type is float, it's stored
				// in fpu registers (st0 through st7). we can use the `fst`
				// instruction to move st0 to a place in memory
				// D9 FSTP m32real ST Store Floating Point Value and Pop
				// D9 = for m32
				callShellcode.push_back(0xD9);
				callShellcode.push_back(0x1C);
				callShellcode.push_back(0x25);
			}
			else if (returnType == Type::T_DOUBLE)
			{
				// fstp QWORD PTR [0x12345678]
				// DD FSTP m64real ST Store Floating Point Value and Pop
				// DD = for m64
				callShellcode.push_back(0xDD);
				callShellcode.push_back(0x1C);
				callShellcode.push_back(0x25);
			}
			else
			{
				// mov [0x1234], eax
				// A3: MOVE moffs16/32 eAX
				// Call routines places return value inside EAX
				callShellcode.push_back(0xA3);
			}

			for (int i = 0; i < 4; i++)
			{
				int shifted = ((DWORD)returnValuePointer >> (i * 8)) & 0xFF;
				callShellcode.push_back(shifted);
			}
		}

		// C3: return
		callShellcode.push_back(0xC3);

		// concatenate the arg shellcode with the calling shellcode
		std::vector<unsigned char> shellcode;
		shellcode.reserve(argShellcode.size() + callShellcode.size());
		shellcode.insert(shellcode.end(), argShellcode.begin(), argShellcode.end());
		shellcode.insert(shellcode.end(), callShellcode.begin(), callShellcode.end());

		// 5 = 0xE8 (call) + 4 empty bytes for where the address will go
		int addessShellcodeOffset = argShellcode.size() + 5;

		// Allocate space for the shellcode
		SIZE_T size = shellcode.size() * sizeof(unsigned char);
		LPVOID pShellcode = VirtualAllocEx(pHandle, NULL, size, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);

		// `call` opcode takes relative address, so calculate the relative address
		// taking into account where the shellcode will be written in memory
		DWORD64 relative = address - (uintptr_t)pShellcode - addessShellcodeOffset;

		// Write the relative address to the shellcode
		for (int i = 0; i < 4; i++)
		{
			int shifted = (relative >> (i * 8)) & 0xFF;

			// argShellcode.size() will offset to the `0xE8` opcode, add 1 to offset that instruction
			shellcode.at(argShellcode.size() + 1 + i) = shifted;
		}

		// Write the shellcode
		WriteProcessMemory(pHandle, pShellcode, shellcode.data(), size, NULL);

		// Execute the shellcode
		HANDLE thread = CreateRemoteThread(pHandle, NULL, NULL, (LPTHREAD_START_ROUTINE)pShellcode, NULL, NULL, NULL);

		Call data = {0, "", (DWORD)-1};

		if (thread == NULL)
		{
			*errorMessage = "unable to create remote thread.";
			return data;
		}

		WaitForSingleObject(thread, INFINITE);
		GetExitCodeThread(thread, &data.exitCode);

		if (returnType != Type::T_VOID && returnType != Type::T_STRING)
		{
			ReadProcessMemory(pHandle, (LPVOID)returnValuePointer, &data.returnValue, sizeof(int), NULL);
			VirtualFreeEx(pHandle, returnValuePointer, 0, MEM_RELEASE);
		}

		if (returnType == Type::T_STRING)
		{
			// String is stored in memory somewhere
			// When returning a string, the address of the string is placed in EAX.
			// So we read the current returnValuePointer address to get the actual address of the string
			ReadProcessMemory(pHandle, (LPVOID)returnValuePointer, &returnValuePointer, sizeof(int), NULL);

			std::vector<char> chars;
			int offset = 0x0;
			while (true)
			{
				char c = functions::readChar(pHandle, (DWORD64)returnValuePointer + offset);
				chars.push_back(c);

				// break at 1 million chars
				if (offset == (sizeof(char) * 1000000))
				{
					chars.clear();
					break;
				}

				// break at terminator (end of string)
				if (c == '\0')
				{
					break;
				}

				// go to next char
				offset += sizeof(char);
			}

			std::string str(chars.begin(), chars.end());
			// TODO: pass str as LPVOID and cast back to string
			data.returnString = str;
		}

		VirtualFreeEx(pHandle, pShellcode, 0, MEM_RELEASE);

		return data;
	}
}

#endif
#pragma once