#include <psapint/psapi.hpp>
#include <common/process.hpp>

ProcessAPI::ProcessAPI(void)
: mpfnEnumProcesses(0), mpfnEnumProcessModules(0), mpfnGetModuleBaseName(0), mpfnGetModuleFileNameEx(0),
	mpfnGetModuleInformation(0), mpfnEmptyWorkingSet(0), mpfnQueryWorkingSet(0), 
	mpfnInitializeProcessForWsWatch(0), mpfnGetWsChanges(0), mpfnGetMappedFileName(0),  
	mpfnEnumDeviceDrivers(0), mpfnGetDeviceDriverBaseName(0), mpfnGetDeviceDriverFileName(0),
	mpfnGetProcessMemoryInfo(0), mPSAPILib("PSAPI.DLL")
{
  getEntryPoints();
}

ProcessAPI::ProcessAPI(const ProcessAPI &someProcessAPI)
{ // private implementation
  *this=someProcessAPI;
}

ProcessAPI::~ProcessAPI()
{
}

ProcessAPI &ProcessAPI::operator=(const ProcessAPI &/*someProcessAPI*/)
{ // private implementation
  return *this;
}

bool ProcessAPI::enumProcesses(ProcessInfoList &processInfoList)
{
  ProcessIDList processIDList;
	String strModuleBaseName;
	String strModuleFileName;

  processInfoList.size(0);
	enumProcesses(processIDList);
	processInfoList.size(processIDList.size());
	for(int procIndex=0;procIndex<processIDList.size();procIndex++)
	{
	  processInfoList[procIndex]=processIDList[procIndex];
		ProcessInfo &processInfo=processInfoList[procIndex];
		enumProcessModules(processInfo.processID(),processInfo);
		ModuleInfoList &moduleInfoList=processInfo;
		for(int modIndex=0;modIndex<moduleInfoList.size();modIndex++)
		{
		  ModuleInfo &moduleInfo=moduleInfoList[modIndex];
		  getModuleBaseName(processInfo.processID(),moduleInfo.module(),strModuleBaseName);
			getModuleFileName(processInfo.processID(),moduleInfo.module(),strModuleFileName);
			moduleInfo.moduleBaseName(strModuleBaseName);
			moduleInfo.moduleFileName(strModuleFileName);
		}
	}
	return processInfoList.size();
}

bool ProcessAPI::enumProcesses(ProcessIDList &processList)
{
  DWORD cbNeeded(0);
	DWORD processListArray[1024];
	DWORD itemCount;

	enumProcesses((unsigned long*)&processListArray,sizeof(processListArray),&cbNeeded);
	itemCount=cbNeeded/sizeof(DWORD);
	processList.size(itemCount);
  for(int index=0;index<itemCount;index++)processList[index]=processListArray[index];
	return processList.size()?true:false;
}

bool ProcessAPI::enumProcessModules(const ProcessID &processID,ModuleInfoList &moduleInfoList)
{
  Process process;
  DWORD cbNeeded(0);
	HMODULE moduleHandleArray[1024];
	DWORD itemCount;

  moduleInfoList.size(0);
	if(!process.openProcess(processID.processID()))return false;
	enumProcessModules(process.getProcess(),(HMODULE*)&moduleHandleArray,sizeof(moduleHandleArray),&cbNeeded);
	itemCount=cbNeeded/sizeof(HANDLE);
  moduleInfoList.size(itemCount);
	for(int index=0;index<itemCount;index++)moduleInfoList[index]=ModuleInfo(moduleHandleArray[index]);
	return moduleInfoList.size()?true:false;
}

DWORD ProcessAPI::getModuleBaseName(const ProcessID &processID,HMODULE hModule,String &strModuleBaseName)
{
  Process process;

	if(!process.openProcess(processID.processID()))return 0;
	strModuleBaseName.reserve(2048);
	return getModuleBaseName(process.getProcess(),hModule,strModuleBaseName,2048);
}

DWORD ProcessAPI::getModuleFileName(const ProcessID &processID,HMODULE hModule,String &strModuleFileName)
{
  Process process;

	if(!process.openProcess(processID.processID()))return 0;
	strModuleFileName.reserve(2048);
	return getModuleFileNameEx(process.getProcess(),hModule,strModuleFileName,2048);
}

bool ProcessAPI::getEntryPoints(void)
{
  if(!mPSAPILib.isOkay())return FALSE;
	mpfnEnumProcesses=(PFNENUMPROCESSES)mPSAPILib.procAddress("EnumProcesses");
	mpfnEnumProcessModules=(PFNENUMPROCESSMODULES)mPSAPILib.procAddress("EnumProcessModules");
  mpfnGetModuleBaseName=(PFNGETMODULEBASENAME)mPSAPILib.procAddress("GetModuleBaseNameA");
	mpfnGetModuleFileNameEx=(PFNGETMODULEFILENAMEEX)mPSAPILib.procAddress("GetModuleFileNameExA");
	mpfnGetModuleInformation=(PFNGETMODULEINFORMATION)mPSAPILib.procAddress("GetModuleInformation");
	mpfnEmptyWorkingSet=(PFNEMPTYWORKINGSET)mPSAPILib.procAddress("EmptyWorkingSet");
	mpfnQueryWorkingSet=(PFNQUERYWORKINGSET)mPSAPILib.procAddress("QueryWorkingSet");
	mpfnInitializeProcessForWsWatch=(PFNINITIALIZEPROCESSFORWSWATCH)mPSAPILib.procAddress("InitializeProcessForWsWatch");
	mpfnGetWsChanges=(PFNGETWSCHANGES)mPSAPILib.procAddress("GetWsChanges");
	mpfnGetMappedFileName=(PFNGETMAPPEDFILENAME)mPSAPILib.procAddress("GetMappedFileNameA");
	mpfnEnumDeviceDrivers=(PFNENUMDEVICEDRIVERS)mPSAPILib.procAddress("EnumDeviceDrivers");
	mpfnGetDeviceDriverBaseName=(PFNGETDEVICEDRIVERBASENAME)mPSAPILib.procAddress("GetDeviceDriverBaseNameA");
	mpfnGetDeviceDriverFileName=(PFNGETDEVICEDRIVERFILENAME)mPSAPILib.procAddress("GetDeviceDriverFileNameA");
	mpfnGetProcessMemoryInfo=(PFNGETPROCESSMEMORYINFO)mPSAPILib.procAddress("GetProcessMemoryInfo");
	return isOkay();
}

bool ProcessAPI::enumProcesses(DWORD *lpidProcess,DWORD cb,DWORD *cbNeeded)
{
  if(!mpfnEnumProcesses){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnEnumProcesses(lpidProcess,cb,cbNeeded);
}

bool ProcessAPI::enumProcessModules(HANDLE hProcess,HMODULE *lphModule,DWORD cb,LPDWORD lpcbNeeded)
{
  if(!mpfnEnumProcessModules){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnEnumProcessModules(hProcess,lphModule,cb,lpcbNeeded);
}

DWORD ProcessAPI::getModuleBaseName(HANDLE hProcess,HMODULE hModule,LPSTR lpBaseName,DWORD nSize)
{
  if(!mpfnGetModuleBaseName){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnGetModuleBaseName(hProcess,hModule,lpBaseName,nSize);
}

DWORD ProcessAPI::getModuleFileNameEx(HANDLE hProcess,HMODULE hModule,LPSTR lpFilename,DWORD nSize)
{
  if(!mpfnGetModuleFileNameEx){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnGetModuleFileNameEx(hProcess,hModule,lpFilename,nSize);
}

bool ProcessAPI::getModuleInformation(HANDLE hProcess,HMODULE hModule,LPMODULEINFO lpmodinfo,DWORD cb)
{
  if(!mpfnGetModuleInformation){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnGetModuleInformation(hProcess,hModule,lpmodinfo,cb);
}

bool ProcessAPI::emptyWorkingSet(HANDLE hProcess)
{
  if(!mpfnEmptyWorkingSet){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnEmptyWorkingSet(hProcess);
}

bool ProcessAPI::queryWorkingSet(HANDLE hProcess,PVOID pv,DWORD cb)
{
  if(!mpfnQueryWorkingSet){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnQueryWorkingSet(hProcess,pv,cb);
}

bool ProcessAPI::initializeProcessForWsWatch(HANDLE hProcess)
{
  if(!mpfnInitializeProcessForWsWatch){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnInitializeProcessForWsWatch(hProcess);
}

bool ProcessAPI::getWsChanges(HANDLE hProcess,PPSAPI_WS_WATCH_INFORMATION lpWatchInfo,DWORD cb)
{
  if(!mpfnGetWsChanges){throw ProcessAPIInvalidEntryPoint();return false;}
	return mpfnGetWsChanges(hProcess,lpWatchInfo,cb);
}

DWORD ProcessAPI::getMappedFileName(HANDLE hProcess,LPVOID lpv,LPWSTR lpFilename,DWORD nSize)
{
  if(!mpfnGetMappedFileName){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnGetMappedFileName(hProcess,lpv,lpFilename,nSize);
}

DWORD ProcessAPI::enumDeviceDrivers(LPVOID *lpImageBase,DWORD cb,LPDWORD lpcbNeeded)
{
  if(!mpfnEnumDeviceDrivers){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnEnumDeviceDrivers(lpImageBase,cb,lpcbNeeded);
}

DWORD ProcessAPI::getDeviceDriverBaseName(LPVOID imageBase,LPSTR lpBaseName,DWORD nSize)
{
  if(!mpfnGetDeviceDriverBaseName){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnGetDeviceDriverBaseName(imageBase,lpBaseName,nSize);
}

DWORD ProcessAPI::getDeviceDriverFileName(LPVOID imageBase,LPSTR lpFilename,DWORD nSize)
{
  if(!mpfnGetDeviceDriverFileName){throw ProcessAPIInvalidEntryPoint();return 0;}
	return mpfnGetDeviceDriverFileName(imageBase,lpFilename,nSize);
}

bool ProcessAPI::getProcessMemoryInfo(HANDLE hProcess,PPROCESS_MEMORY_COUNTERS ppsmemCounters,DWORD cb)
{
  if(!mpfnGetProcessMemoryInfo){throw ProcessAPIInvalidEntryPoint();return false;}  
	return mpfnGetProcessMemoryInfo(hProcess,ppsmemCounters,cb);
}

bool ProcessAPI::isOkay(void)const
{
  if(!mPSAPILib.isOkay()||!mpfnEnumProcesses||!mpfnEnumProcessModules||!mpfnGetModuleBaseName||
	   !mpfnGetModuleFileNameEx||!mpfnGetModuleInformation||!mpfnEmptyWorkingSet||!mpfnQueryWorkingSet||
	   !mpfnInitializeProcessForWsWatch||!mpfnGetWsChanges||!mpfnGetMappedFileName||!mpfnEnumDeviceDrivers||
		 !mpfnGetDeviceDriverBaseName||!mpfnGetDeviceDriverFileName||!mpfnGetProcessMemoryInfo)return FALSE;
	return TRUE;
}