/* DLL injector
 * Author: gwangyi (gwangyi@postech.ac.kr)
 *
 * This program is available under the Creative Commons-Attribution-Noncommercial-Share Alike 2.0 Korea license or any later
 */
#include <windows.h>
#include <stdio.h>
#include <tchar.h>
#include <psapi.h>

HINSTANCE g_hInst = NULL;

LRESULT WINAPI MainWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam);

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInst, LPSTR lpszCmdLine, int nCmdShow)
{
	HWND hWnd;
	WNDCLASS wndclass;
	MSG msg;
	
	g_hInst = hInstance;
	
	ZeroMemory(&wndclass, sizeof(wndclass));
	wndclass.hCursor = LoadCursor(NULL, IDC_ARROW);
	wndclass.hIcon = LoadIcon(NULL, IDI_APPLICATION);
	wndclass.hInstance = hInstance;
	wndclass.hbrBackground = (HBRUSH)GetStockObject(WHITE_BRUSH);
	wndclass.lpfnWndProc = MainWndProc;
	wndclass.lpszClassName = TEXT("DLLInjector");
	wndclass.style = CS_HREDRAW | CS_VREDRAW;
	RegisterClass(&wndclass);
	
	hWnd = CreateWindow(TEXT("DLLInjector"), TEXT("DLL Injector"), WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, NULL, NULL, hInstance, NULL);
	ShowWindow(hWnd, nCmdShow);
	UpdateWindow(hWnd);
	
	while(GetMessage(&msg, NULL, 0, 0))
	{
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	return (int)msg.wParam;
}

BOOL FillProcessList(HWND hWnd)
{
	DWORD aProcesses[1024], cbNeeded, cProcesses;

	unsigned int i;
	
	if(!EnumProcesses(aProcesses, sizeof(aProcesses), &cbNeeded))
		return FALSE;
	
	cProcesses = cbNeeded / sizeof(DWORD);
	for(i = 0; i < cProcesses; i++)
	{
	    TCHAR szProcessName[MAX_PATH] = TEXT("[unknown]");
	    TCHAR item[MAX_PATH + 10] = {0, };
	    LRESULT idx;

	    // Get a handle to the process.

	    HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
	                                   PROCESS_VM_READ,
	                                   FALSE, aProcesses[i] );

	    // Get the process name.

	    if (NULL != hProcess )
	    {
	        HMODULE hMod;
	        DWORD cbNeeded;

	        if ( EnumProcessModules( hProcess, &hMod, sizeof(hMod), 
	             &cbNeeded) )
	        {
	            GetModuleBaseName( hProcess, hMod, szProcessName, 
	                               sizeof(szProcessName)/sizeof(TCHAR) );
	        }
	    }

	    // Print the process name and identifier.

	    _stprintf( item, TEXT("%s (PID: %lu)"), szProcessName, aProcesses[i] );
	    idx = SendDlgItemMessage(hWnd, 1, LB_ADDSTRING, 0, (LPARAM)item);
	    if(idx == LB_ERR || idx == LB_ERRSPACE) return FALSE;
	    SendDlgItemMessage(hWnd, 1, LB_SETITEMDATA, idx, aProcesses[i]);

	    CloseHandle( hProcess );
	}
	return TRUE;
}

LRESULT WINAPI MainWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	switch(uMsg)
	{
	case WM_CREATE:
		CreateWindow(TEXT("LISTBOX"), TEXT(""), WS_BORDER | WS_CHILD | WS_VISIBLE | LBS_SORT | WS_VSCROLL, 0, 0, 100, 100, hWnd, (HMENU)1, g_hInst, NULL);
		CreateWindow(TEXT("BUTTON"), TEXT("&Refresh"), WS_CHILD | WS_VISIBLE, 0, 0, 100, 100, hWnd, (HMENU)2, g_hInst, NULL);
		CreateWindow(TEXT("EDIT"), TEXT(""), WS_BORDER | WS_CHILD | WS_VISIBLE, 0, 0, 100, 100, hWnd, (HMENU)3, g_hInst, NULL);
		CreateWindow(TEXT("BUTTON"), TEXT("&..."), WS_CHILD | WS_VISIBLE, 0, 0, 100, 100, hWnd, (HMENU)4, g_hInst, NULL);
		CreateWindow(TEXT("BUTTON"), TEXT("&Inject!"), WS_CHILD | WS_VISIBLE, 0, 0, 100, 100, hWnd, (HMENU)5, g_hInst, NULL);
		if(!FillProcessList(hWnd)) return -1;
		break;
	case WM_SIZE:
	{
		RECT rect;
		GetClientRect(hWnd, &rect);
		if(rect.bottom - rect.top <= 0 || rect.right - rect.left <= 0) break;
		SetWindowPos(GetDlgItem(hWnd, 1), NULL, 10, 10, rect.right - rect.left - 20, rect.bottom - rect.top - 50, SWP_NOZORDER);
		SetWindowPos(GetDlgItem(hWnd, 2), NULL, 10, rect.bottom - rect.top - 40, 80, 30, SWP_NOZORDER);
		SetWindowPos(GetDlgItem(hWnd, 3), NULL, 100, rect.bottom - rect.top - 40, rect.right - rect.left - 290, 30, SWP_NOZORDER);
		SetWindowPos(GetDlgItem(hWnd, 4), NULL, rect.right - 180, rect.bottom - rect.top - 40, 80, 30, SWP_NOZORDER);
		SetWindowPos(GetDlgItem(hWnd, 5), NULL, rect.right - 90, rect.bottom - rect.top - 40, 80, 30, SWP_NOZORDER);
	}
		break;
	case WM_COMMAND:
		switch(LOWORD(wParam))
		{
		case 2:
			if(HIWORD(wParam) == 0)
			{
				SendDlgItemMessage(hWnd, 1, LB_RESETCONTENT, 0, 0);
				FillProcessList(hWnd);
			}
			break;
		case 4:
			if(HIWORD(wParam) == 0)
			{
				OPENFILENAME ofn;
				TCHAR szFile[MAX_PATH];
				
				GetDlgItemText(hWnd, 3, szFile, MAX_PATH);
				ZeroMemory(&ofn, sizeof(ofn));
				ofn.lStructSize = sizeof(ofn);
				ofn.hwndOwner = hWnd;
				ofn.lpstrFile = szFile;
				ofn.nMaxFile = MAX_PATH;
				ofn.lpstrFilter = TEXT("Dynamic Linked Library\0*.dll\0");
				ofn.nFilterIndex = 0;
				ofn.lpstrFileTitle = NULL;
				ofn.nMaxFileTitle = 0;
				ofn.lpstrInitialDir = NULL;
				ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;
				
				if(GetOpenFileName(&ofn) == TRUE)
				{
					SetDlgItemText(hWnd, 3, szFile);
				}
			}
			break;
		case 5:
			if(HIWORD(wParam) == 0)
			{
				char szDllName[MAX_PATH];
				DWORD dwPID, dwRet;
				HANDLE hProc, hThread;
				HANDLE hKernel32;
				LRESULT ret = SendDlgItemMessage(hWnd, 1, LB_GETCURSEL, 0, 0);
				LPTHREAD_START_ROUTINE pfnLoadLibrary;
				LPVOID pRemoteDll;

				if(ret == LB_ERR) break;
				dwPID = SendDlgItemMessage(hWnd, 1, LB_GETITEMDATA, ret, 0);
				GetDlgItemTextA(hWnd, 3, szDllName, MAX_PATH);
				
				hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID);
				if(hProc == NULL)
				{
					MessageBox(hWnd, TEXT("Failed to inject DLL!\nCannot open the process"), NULL, MB_ICONERROR);
					break;
				}
				pRemoteDll = VirtualAllocEx(hProc, NULL, sizeof(char) * MAX_PATH, MEM_COMMIT, PAGE_READWRITE);
				if(!pRemoteDll)
				{
					CloseHandle(hProc);
					MessageBox(hWnd, TEXT("Failed to inject DLL!\nCannot allocate memory in target process"), NULL, MB_ICONERROR);
					break;
				}
				if(!WriteProcessMemory(hProc, pRemoteDll, szDllName, sizeof(szDllName), NULL))
				{
					VirtualFreeEx(hProc, pRemoteDll, 0, MEM_RELEASE);
					CloseHandle(hProc);
					MessageBox(hWnd, TEXT("Failed to inject DLL!\nCannot write memory in target process"), NULL, MB_ICONERROR);
					break;
				}
				hKernel32 = LoadLibrary(TEXT("KERNEL32.DLL"));
				pfnLoadLibrary = (LPTHREAD_START_ROUTINE)GetProcAddress(hKernel32, TEXT("LoadLibraryA"));
				
				hThread = CreateRemoteThread(hProc, NULL, 0, pfnLoadLibrary, pRemoteDll, 0, NULL);
				if(!hThread)
				{
					VirtualFreeEx(hProc, pRemoteDll, 0, MEM_RELEASE);
					CloseHandle(hProc);
					FreeLibrary(hKernel32);
					MessageBox(hWnd, TEXT("Failed to inject DLL!\nCannot create remote thread"), NULL, MB_ICONERROR);
					break;
				}
				WaitForSingleObject(hThread, INFINITE);
				
				VirtualFreeEx(hProc, pRemoteDll, 0, MEM_RELEASE);
				CloseHandle(hProc);
				FreeLibrary(hKernel32);
				if(GetExitCodeThread(hThread, &dwRet))
				{
					TCHAR out[1024];
					_stprintf(out, TEXT("Injection finished.\nExit code: 0x%08lX"), dwRet);
					MessageBox(hWnd, out, TEXT("Finished"), MB_ICONINFORMATION);
				}
			}
			break;
		}
		break;
	case WM_DESTROY:
		PostQuitMessage(0);
		break;
	}
	return DefWindowProc(hWnd, uMsg, wParam, lParam);
}