/*****************************************************************************
 * Inject - Inject code into a remote process.                               *
 *          InjectDll()/EjectDll() - Inject a library into a remote process  *
 *          using the CreateRemoteThread()/LoadLibrary() technique.          *
 *          InjectCode()/EjectCode() - Inject code into a remote process     *
 *          using the CreateRemoteThread()/WriteProcessMemory() technique.   *
 *                                                                           *
 * (c) A. Miguel Feij鉶, 1/6/2004                                            *
 *                                                                           *
 * Adapted from the original code by Robert Kuster.                          *
 *                                                                           *
 *****************************************************************************/

#include <windows.h>
#include <tlhelp32.h>
#include "Inject.h"

// Remove comment if you want to inject the library
//#define INJECT_DLL 1

/**************************************
 * Inject code into a remote process. *
 **************************************/
BOOL Inject()
{
#ifdef INJECT_DLL
	return (InjectDll() != 0);
#else
	return (InjectCode() != 0);
#endif
}


/*************************************
 * Eject code from a remote process. *
 *************************************/
BOOL Eject()
{
#ifdef INJECT_DLL
	return (EjectDll() != 0);
#else
	return (EjectCode() != 0);
#endif
}

////////////////////////////////// General purpose routines //////////////////////

/******************************
 * Enable/Disable privilege.  *
 * Called with SE_DEBUG_NAME. *
 ******************************/
BOOL EnablePrivilege(LPCTSTR lpszPrivilegeName, BOOL bEnable)
{
	HANDLE				hToken;
	TOKEN_PRIVILEGES	tp;
	LUID				luid;
	BOOL				ret;

	if (!OpenProcessToken(GetCurrentProcess(),
					      TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY | TOKEN_READ,
					      &hToken))
		return FALSE;

	if (!LookupPrivilegeValue(NULL, lpszPrivilegeName, &luid))
		return FALSE;

	tp.PrivilegeCount           = 1;
	tp.Privileges[0].Luid       = luid;
	tp.Privileges[0].Attributes = bEnable ? SE_PRIVILEGE_ENABLED : 0;

	ret = AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL);

	CloseHandle(hToken);

	return ret;
}


/********************************************
 * Return PID using the Toolhelp functions. *
 ********************************************/
DWORD GetPIDFromNameToolhelp(char *szProcessName)
{
    typedef HANDLE (WINAPI *CREATESNAPSHOT) (DWORD, DWORD);
    typedef BOOL   (WINAPI *PROCESSWALK)    (HANDLE, LPPROCESSENTRY32); 

    HINSTANCE       hKernel;
    CREATESNAPSHOT  CreateToolhelp32Snapshot;
    PROCESSWALK     Process32First;
    PROCESSWALK     Process32Next;

    HANDLE          hSnapshot;
    PROCESSENTRY32  pe32;
    BOOL            bRes;
    char            *p;
    DWORD           dwPID = -1;

	// Check szProcessName
	if (!szProcessName)
		return -1;

    // Get Kernel32 handle
    if (!(hKernel = GetModuleHandle("Kernel32.dll")))
        return -1;

    // We must link to these functions explicitly.
    // Otherwise it will fail on Windows NT which doesn't have Toolhelp
    // functions defined in Kernel32.
    CreateToolhelp32Snapshot = (CREATESNAPSHOT) GetProcAddress(hKernel, "CreateToolhelp32Snapshot");
    Process32First = (PROCESSWALK) GetProcAddress(hKernel, "Process32First");
    Process32Next = (PROCESSWALK) GetProcAddress(hKernel, "Process32Next");
    if (!CreateToolhelp32Snapshot || !Process32First || !Process32Next)
	{
		FreeLibrary(hKernel);
		SetLastError(ERROR_PROC_NOT_FOUND);
        return -1;
	}

    hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return -1;

    pe32.dwSize = sizeof(pe32);
    bRes = Process32First(hSnapshot, &pe32);
    
    while (bRes)
    {
        // Strip off full path
        p = strrchr(pe32.szExeFile, '\\');
        if (p)
            p++;
        else
            p = pe32.szExeFile;

        // Process found ?
        if (stricmp(p, szProcessName) == 0)
        {
            dwPID = pe32.th32ProcessID;
            break;
        }

        bRes = Process32Next(hSnapshot, &pe32);
    }

    CloseHandle(hSnapshot);
    return dwPID;
}


/*****************************************
 * Return PID using the PSAPI functions. *
 *****************************************/
DWORD GetPIDFromNamePSAPI(char *szProcessName)
{
    typedef BOOL  (WINAPI *ENUMPROCESSES) (DWORD *, DWORD, DWORD *);
    typedef BOOL  (WINAPI *ENUMPROCESSMODULES) (HANDLE, HMODULE *, DWORD, LPDWORD);
    typedef DWORD (WINAPI *GETMODULEBASENAME) (HANDLE, HMODULE, LPTSTR, DWORD);

	DWORD				dwPID = -1;				
    HINSTANCE           hPsApi;
    ENUMPROCESSES       EnumProcesses;
    ENUMPROCESSMODULES  EnumProcessModules;
	GETMODULEBASENAME	GetModuleBaseName;
    DWORD               * PIDArray = NULL;
	DWORD               cbAlloc = 128;
    DWORD               cbNeeded;
    DWORD               nProcesses;
    HANDLE              hProcess;
    DWORD               i;
    HMODULE             hExeModule;
    char                szModuleName[MAX_PATH];

	// Check szProcessName
	if (!szProcessName)
		return -1;

    // Get PSAPI handle
    if (!(hPsApi = LoadLibrary("PSAPI.DLL")))
        return -1;

    // Load functions in PSAPI dynamically.
    // We can't be sure PSAPI.DLL is installed.
    EnumProcesses = (ENUMPROCESSES) GetProcAddress(hPsApi, "EnumProcesses");
    EnumProcessModules = (ENUMPROCESSMODULES) GetProcAddress(hPsApi, "EnumProcessModules");
	GetModuleBaseName = (GETMODULEBASENAME) GetProcAddress(hPsApi, "GetModuleBaseNameA");
    if (!EnumProcesses || !EnumProcessModules || !GetModuleBaseName)
	{
		FreeLibrary(hPsApi);
		SetLastError(ERROR_PROC_NOT_FOUND);
        return -1;
	}

    // Find needed buffer length
	do
	{
	    cbAlloc *= 2;

		if (PIDArray)
			free(PIDArray);

		PIDArray = malloc(cbAlloc);
		if (!PIDArray)
		{
			FreeLibrary(hPsApi);
			SetLastError(ERROR_NOT_ENOUGH_MEMORY);
	        return -1;
		}
		
		if (!EnumProcesses(PIDArray, cbAlloc, &cbNeeded))
		{
			FreeLibrary(hPsApi);
			free(PIDArray);
	        return -1;
		}

	} while (cbNeeded == cbAlloc);

    nProcesses = cbNeeded / sizeof(DWORD);

    for (i=0; i < nProcesses; i++)
    {
        hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
                               FALSE,
                               PIDArray[i]);
        if (!hProcess)
            continue;

        if (!EnumProcessModules(hProcess, &hExeModule, sizeof(hExeModule), &cbNeeded))
        {
            CloseHandle(hProcess);
            continue;
        }

        GetModuleBaseName(hProcess,
                          hExeModule,
                          szModuleName,
                          sizeof(szModuleName));

        // Process found ?
        if (stricmp(szModuleName, szProcessName) == 0)
        {
            dwPID = PIDArray[i];
			CloseHandle(hProcess);
            break;
        }

        CloseHandle(hProcess);
    }

	FreeLibrary(hPsApi);
	free(PIDArray);
    return dwPID;
}


/************************************************
 * Return PID using ZwQuerySystemInformation(). *
 ************************************************/
DWORD GetPIDFromNameZwQuerySystemInformation(char *szProcessName)
{
	// Some definitions from NTDDK and other sources

	typedef LONG	NTSTATUS;
	typedef LONG	KPRIORITY;

	#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >= 0)
	#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)
	#define SystemProcessesAndThreadsInformation	5

	typedef struct _CLIENT_ID {
		DWORD	    UniqueProcess;
		DWORD	    UniqueThread;
	} CLIENT_ID;

	typedef struct _UNICODE_STRING {
		USHORT	    Length;
		USHORT	    MaximumLength;
		PWSTR	    Buffer;
	} UNICODE_STRING;

	typedef struct _VM_COUNTERS {
		SIZE_T	    PeakVirtualSize;
		SIZE_T	    VirtualSize;
		ULONG	    PageFaultCount;
		SIZE_T	    PeakWorkingSetSize;
		SIZE_T	    WorkingSetSize;
		SIZE_T	    QuotaPeakPagedPoolUsage;
		SIZE_T	    QuotaPagedPoolUsage;
		SIZE_T	    QuotaPeakNonPagedPoolUsage;
		SIZE_T	    QuotaNonPagedPoolUsage;
		SIZE_T	    PagefileUsage;
		SIZE_T	    PeakPagefileUsage;
	} VM_COUNTERS;

	typedef struct _SYSTEM_THREADS {
		LARGE_INTEGER   KernelTime;
		LARGE_INTEGER   UserTime;
		LARGE_INTEGER   CreateTime;
		ULONG			WaitTime;
		PVOID			StartAddress;
		CLIENT_ID	    ClientId;
		KPRIORITY	    Priority;
		KPRIORITY	    BasePriority;
		ULONG			ContextSwitchCount;
		LONG			State;
		LONG			WaitReason;
	} SYSTEM_THREADS, * PSYSTEM_THREADS;

	// NOTE: SYSTEM_PROCESSES structure is different on NT 4 and Win2K
	typedef struct _SYSTEM_PROCESSES {
		ULONG			NextEntryDelta;
		ULONG			ThreadCount;
		ULONG			Reserved1[6];
		LARGE_INTEGER   CreateTime;
		LARGE_INTEGER   UserTime;
		LARGE_INTEGER   KernelTime;
		UNICODE_STRING  ProcessName;
		KPRIORITY	    BasePriority;
		ULONG			ProcessId;
		ULONG			InheritedFromProcessId;
		ULONG			HandleCount;
		ULONG			Reserved2[2];
		VM_COUNTERS	    VmCounters;
	#if _WIN32_WINNT >= 0x500
		IO_COUNTERS	    IoCounters;
	#endif
		SYSTEM_THREADS  Threads[1];
	} SYSTEM_PROCESSES, * PSYSTEM_PROCESSES;

	typedef LONG (WINAPI *ZWQUERYSYSTEMINFORMATION)(UINT SystemInformationClass, PVOID SystemInformation, ULONG SystemInformationLength, PULONG ReturnLength);

	ZWQUERYSYSTEMINFORMATION    ZwQuerySystemInformation;
	PSYSTEM_PROCESSES           pInfo;
    HINSTANCE					hNTDll;
	WCHAR						wcProcessName[MAX_PATH+2];
	PCWSTR						pszProcessName;
	DWORD						dwPID;
	ULONG						BufferLen = 0x8000;
	LPVOID						pBuffer = NULL;
	LONG						Status;

	// Check szProcessName
	if (!szProcessName)
		return -1;

	// Convert szProcessName to Unicode
	if (!MultiByteToWideChar(CP_ACP, 0, szProcessName, -1, wcProcessName, sizeof(wcProcessName) / sizeof(WCHAR)))
		return -1;

    // Get NTDLL handle