#ifdef _MSC_VER
#define _CRT_SECURE_NO_WARNINGS
#endif
#include <Windows.h>
#include <TlHelp32.h>
#include <Psapi.h>
#include <iostream>
#include <stdlib.h>

#define DUMP "C:\\temp\\dump.bin"

class VictimProcess
{
public:
 
    PROCESSENTRY32 __gameProcess;
	DWORD64 __dwordProcess;
    HANDLE __HandleProcess;
    HWND __HWNDCsgo;
	HMODULE __HModuleTarget;
	LPVOID __dwordTarget;
	LPSTR __sBinary;
	LPSTR __sCmdLine;
	LPSTR __sProcess;
	LPSTR __sDumpPath;
	LPSTR __sDumpFile;
 
    void RunProcess()
    {
		PROCESS_INFORMATION lpi;
		STARTUPINFO si;
		memset(&lpi,0,sizeof(PROCESS_INFORMATION));
		memset(&si,0,sizeof(STARTUPINFO));

		std::cout<<"Running target binary"<<std::endl;

		if(!CreateProcess(NULL,__sCmdLine,NULL,NULL,FALSE,(CREATE_NO_WINDOW),NULL,NULL,&si,&lpi))	//Runs victim process
		{
			quit("Couldn't start process, press something to quit");
		}

		__HandleProcess = lpi.hProcess;			//Handle to victim process
		runSetDebugPrivs();						//Sets debug privs on victim process, allows for dumping of memory
    }

	void dumpModule(LPSTR sTarget)
	{
		while(__dwordTarget == 0x0) __dwordTarget = GetModuleNamePointer(sTarget, __gameProcess.th32ProcessID);	//Walks through processes modules until name match

		memset(&__HModuleTarget,0,sizeof(HMODULE));
		__HModuleTarget = (HMODULE)__dwordTarget;
		MODULEINFO modInfo = getModInfo(__HModuleTarget);			//Grabs module info for target module

		printf("Address of %s: %X \n",sTarget,__dwordTarget);
		printf("Size of %s: %d \n",sTarget,modInfo.SizeOfImage);

		initDump(sTarget);											//Sets up dump file

		walkMod(modInfo.lpBaseOfDll,modInfo.SizeOfImage);			//Walks and dumps module to disk
	}

	void findProcess(LPCSTR sTarget)	
	{
		while(!FindProcessName(sTarget,&__gameProcess)) Sleep(10);								//Finds process by name
		while(!getThreadByProcess(__gameProcess.th32ProcessID)) Sleep(10);						//Grabs main threadId
		__HandleProcess = OpenProcess(PROCESS_ALL_ACCESS, false, __gameProcess.th32ProcessID);	//Grabs handle to process
	}

	void killPProcess()			//Kills handle of parent calling process
	{
		if(__HandleProcess == NULL)
		{
			quit("Can't get process name, exiting...");
		}
		if(!TerminateProcess(__HandleProcess,0))
			quit("Unable to terminate process");
		quit("Things seem to have gone smoothly...");
	}

	VictimProcess(LPSTR sBinary, LPSTR sCmd, LPSTR sDumpPath)
	{
		__sBinary = sBinary;
		char buff[_MAX_PATH+100];

		memset(buff,0,_MAX_PATH+100);
		strcpy_s(buff,_MAX_PATH+100,sBinary);	//Copies binary path, and cmd args to the same buffer
		strcat_s(buff,_MAX_PATH+100," ");
		strcat_s(buff,_MAX_PATH+100,sCmd);

		size_t len = strlen(buff)+1;
		__sCmdLine = new char[len];				
		memset(__sCmdLine, 0, len);
		strcpy_s(__sCmdLine,len, buff);			//Copies cmd line call to __sCmdLine

		char drive[_MAX_DRIVE];
		char dir[_MAX_DIR];
		char fname[_MAX_FNAME];
		char ext[_MAX_EXT];

		_splitpath_s(sBinary,drive,_MAX_DRIVE,dir,_MAX_DIR,fname,_MAX_FNAME,ext,_MAX_EXT);	//Splits the file path into different portions

		memset(buff,0,_MAX_PATH+100);
		strcpy_s(buff,_MAX_PATH+100,fname);
		strcat_s(buff,_MAX_PATH+100,ext);
		len = strlen(buff)+1;
		__sProcess = new char[len];				//Gets filename of binary, which is name of process
		memset(__sProcess, 0, len);
		strcpy_s(__sProcess,len, buff);

		len = strlen(sDumpPath)+2;
		__sDumpPath = new char[len];			//Sets dump path
		memset(__sDumpPath,0,len);
		strcpy_s(__sDumpPath,len,sDumpPath);

		CreateDirectory(sDumpPath,NULL);
	}

private:

    MODULEINFO getModInfo(HMODULE hModule)
	{
		MODULEINFO modInfo;
		memset(&modInfo,0,sizeof(MODULEINFO));
		if(!GetModuleInformation(__HandleProcess,hModule,&modInfo,sizeof(MODULEINFO)))		//Grabs module info 
		{
			DWORD x = GetLastError();
			char msg[50];
			sprintf(msg,"GetModuleInformation failed.... Error: %X\n",x);
			quit(msg);
		}

		return modInfo;
	}

	LPVOID GetModuleNamePointer(LPSTR LPSTRModuleName, DWORD __DwordProcessId)
    {
		MODULEENTRY32 lpModuleEntry = {0};
		HANDLE hSnapShot = CreateToolhelp32Snapshot( TH32CS_SNAPMODULE, __DwordProcessId);
		if(!hSnapShot)
			return NULL; 
		lpModuleEntry.dwSize = sizeof(lpModuleEntry);

		BOOL __RunModule = Module32First( hSnapShot, &lpModuleEntry );						//Grabs first module in process list

		
		do{if(!_strcmpi(lpModuleEntry.szModule, LPSTRModuleName ) )							//Checks if module name matches target
			{
				CloseHandle( hSnapShot );
				return (LPVOID)lpModuleEntry.modBaseAddr;									//Returns base address of module
			}
		}while(Module32Next( hSnapShot, &lpModuleEntry ));									//Grabs the next module on the list

		CloseHandle( hSnapShot );
		return NULL;
    }

	DWORD FindProcessName(const char *__ProcessName, PROCESSENTRY32 *pEntry)
    {   
        PROCESSENTRY32 __ProcessEntry;
        __ProcessEntry.dwSize = sizeof(PROCESSENTRY32);
        HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (hSnapshot == INVALID_HANDLE_VALUE) return 0;

		if (!Process32First(hSnapshot, &__ProcessEntry))								//Grabs first process in OS list
        {
            CloseHandle(hSnapshot);
            return 0;
        }
        do{if (!_strcmpi(__ProcessEntry.szExeFile, __ProcessName))						//Checks if process name matches target
        {
            memcpy((void *)pEntry, (void *)&__ProcessEntry, sizeof(PROCESSENTRY32));	//Stores process entry
            CloseHandle(hSnapshot);
            return __ProcessEntry.th32ProcessID;										//Returns target PID
        }} while (Process32Next(hSnapshot, &__ProcessEntry));							//Grabs next process in list
        CloseHandle(hSnapshot);
        return 0;
    }
 
    DWORD getThreadByProcess(DWORD __DwordProcess)
    {   
            THREADENTRY32 __ThreadEntry;
            __ThreadEntry.dwSize = sizeof(THREADENTRY32);
            HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
            if (hSnapshot == INVALID_HANDLE_VALUE) return 0;
 
            if (!Thread32First(hSnapshot, &__ThreadEntry)) {CloseHandle(hSnapshot); return 0; }	//Grabs first threadEntry
 
            do {if (__ThreadEntry.th32OwnerProcessID == __DwordProcess)							//Does threadID match PID
            {
                CloseHandle(hSnapshot);
                return __ThreadEntry.th32ThreadID;												//Returns Thread ID of process
            }} while (Thread32Next(hSnapshot, &__ThreadEntry));									//Grabs next ThreadEntry off list
            CloseHandle(hSnapshot);      
            return 0;
    }
 
	//Sets up Debug Privileges for Current Process Token
    void runSetDebugPrivs()
    {
        HANDLE HandleProcess=GetCurrentProcess(), __HandleToken;
        TOKEN_PRIVILEGES priv;
        LUID __LUID;

        OpenProcessToken(HandleProcess, TOKEN_ADJUST_PRIVILEGES, &__HandleToken);	//Opens accessor for current process
        LookupPrivilegeValue(0, "seDebugPrivilege", &__LUID);						//Gets value for Debug Privs

        priv.PrivilegeCount = 1;
        priv.Privileges[0].Luid = __LUID;											//Sets token values for debug privs
        priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

        AdjustTokenPrivileges(__HandleToken, false, &priv, 0, 0, 0);				//Applies debug token to current process
        CloseHandle(__HandleToken);
        CloseHandle(HandleProcess);
    }
     
	//Walks pages of process space and dumps memory
	DWORD64 walkMod(LPVOID baseAddr,DWORD modSz)
	{
		DWORD64 memoryCount = 0;
		BYTE* p = new BYTE[modSz];
		SIZE_T read;

		ReadProcessMemory(__HandleProcess,(LPCVOID)__dwordTarget,p,(SIZE_T)modSz,&read);	//Grabs chunk of memory from module address space

		printf_s("Read %X bytes of %X reported\n",read,modSz);
		if(!writeDump(p,modSz))																//Write module chunk to disk
		{
			std::cout<<"Write to disk failed"<<std::endl;
			return -1;
		}

		return read;
	}


	//Initiates file to be dumped to
	DWORD64 initDump(LPSTR sTarget)
	{
		if(__sDumpFile != NULL)
			remove(__sDumpFile);

		size_t len = strlen(sTarget)+strlen(__sDumpPath) + 10;
		__sDumpFile = new char[len];
		strcpy_s(__sDumpFile,len,__sDumpPath);			//Builds dump file path
		strcat_s(__sDumpFile,len,sTarget);
		strcat_s(__sDumpFile,len,".bin");

		FILE* fPtr = fopen(__sDumpFile, "wb");
		rewind(fPtr);
		fclose(fPtr);

		return 0;
	}

	//Appends memory chunk to file
	bool writeDump(BYTE* p,SIZE_T sz)
	{
		FILE* fPtr = fopen(__sDumpFile, "ab");
		DWORD64 ret = fwrite(p,sz,1,fPtr);
		fclose(fPtr);

		if(ret < 1)
			return false;
		else
			return true;
	}

	void quit(LPSTR s)
	{
		std::cout<<s<<std::endl;
		exit(0);
	}

};
 
extern VictimProcess fProcess;