#include "stdafx.h"
#include "Injector.h"
#include <tlhelp32.h>


HMODULE k_Injector::TryInjectDll(DWORD adw_ProcessId, const std::wstring& as_DllFile)
{
    //Find the address of the LoadLibrary api, luckily for us, it is loaded in the same address for every process
    HMODULE hLocKernel32 = GetModuleHandleW(L"KERNEL32");
    FARPROC hLocLoadLibrary = GetProcAddress(hLocKernel32, "LoadLibraryW");
    
    //Adjust token privileges to open system processes
    HANDLE hToken;
    TOKEN_PRIVILEGES tkp;
    if(OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
    {
        LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tkp.Privileges[0].Luid);
        tkp.PrivilegeCount = 1;
        tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        AdjustTokenPrivileges(hToken, 0, &tkp, sizeof(tkp), NULL, NULL);
		CloseHandle(hToken);
    }

    //Open the process with all access
    HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, adw_ProcessId);
	if (hProc == NULL)
		return false;

    //Allocate memory to hold the path to the Dll File in the process's memory
    LPVOID hRemoteMem = VirtualAllocEx(hProc, NULL, as_DllFile.size()*sizeof(wchar_t), MEM_COMMIT, PAGE_READWRITE);

    //Write the path to the Dll File in the location just created
    DWORD numBytesWritten;
    WriteProcessMemory(hProc, hRemoteMem, as_DllFile.c_str(), as_DllFile.size()*sizeof(wchar_t), &numBytesWritten);

    //Create a remote thread that starts begins at the LoadLibrary function and is passed are memory pointer
    HANDLE hRemoteThread = CreateRemoteThread(hProc, NULL, 0, (LPTHREAD_START_ROUTINE)hLocLoadLibrary, hRemoteMem, 0, NULL);

	//Wait for the thread to finish
	::WaitForSingleObject( hRemoteThread, INFINITE );
	DWORD  hLibModule = 0;
    ::GetExitCodeThread( hRemoteThread, &hLibModule );

    //Free the memory created on the other process
	::VirtualFreeEx(hProc, hRemoteMem, as_DllFile.size()*sizeof(wchar_t), MEM_RELEASE);

    //Release the handle to the other process
	::CloseHandle(hProc);

    return (HMODULE)hLibModule;
}


bool k_Injector::TryUnInjectDll(DWORD adw_ProcessId, HMODULE ah_ModuleHandle)
{
	//Open the process with all access
    HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, adw_ProcessId);
	if (hProc == NULL)
		return false;

	bool lb_ReturnValue = false;

	HMODULE hLocKernel32 = GetModuleHandleW(L"KERNEL32");
    FARPROC hLocLoadLibrary = GetProcAddress(hLocKernel32, "FreeLibrary");

    if(ah_ModuleHandle != NULL)
    {
        HANDLE hRemoteThread = ::CreateRemoteThread(hProc, NULL, 0, (LPTHREAD_START_ROUTINE)hLocLoadLibrary, (void*)ah_ModuleHandle, 0, NULL );

        if( hRemoteThread != NULL )
        {
			DWORD ldw_ReturnCode;
            ::WaitForSingleObject( hRemoteThread, INFINITE );
            ::GetExitCodeThread( hRemoteThread, &ldw_ReturnCode );
            ::CloseHandle( hRemoteThread );

			lb_ReturnValue = ldw_ReturnCode != 0;
        }
    }

	::CloseHandle(hProc);

	return lb_ReturnValue;
}
