#include "pslib.h"
#include "lprintf.h"
#include <stdio.h>
#include <string.h>
#define DMID 5
#include "debug.h"

#ifndef WIN32
#	define stricmp strcasecmp
#endif

#ifndef MAX_PATH
#	define MAX_PATH 256
#endif

typedef struct find_data {
	char *filename;
	int pid;
} fData;

typedef struct lower_priority_data {
	char *process;
	int retcode;
} lPdata;

/* Library handle */
static HINSTANCE hInstLib;
static int initialized;

/* Function pointers */
static BOOL (WINAPI *lpfEnumProcesses)(DWORD *, DWORD, DWORD *);
static BOOL (WINAPI *lpfEnumProcessModules)(HANDLE, HMODULE *, DWORD, LPDWORD);
static DWORD (WINAPI *lpfGetModuleFileNameEx)(HANDLE, HMODULE, LPTSTR, DWORD);

/* Internal functions */
static char system_error_string[101];
static void pslib_set_error(void) {
	FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM|FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		GetLastError(),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPTSTR) system_error_string,
		100,
		NULL
	);
}

static int find_callbk(void *xdata, char *filename, int pid) {
	fData *data = (fData *)xdata;

	if (stricmp(filename,data->filename)==0) {
		data->pid = pid;
		return FALSE;
	}
	return TRUE;
}

static int pslib_cb_lower(void *xdata, char *filename, int pid) {
	lPdata *data = (lPdata *)xdata;

	if (stricmp(filename,data->process)==0) {
		data->retcode = pslib_priority(pid,BELOWNORMAL_PRIORITY_CLASS);
		return FALSE;
	}
	return TRUE;
}


static void strip_path(char *filename) {
	char *ptr;

	for (ptr = filename + (int)strlen(filename) - 1; TRUE; ptr--) {
		if (*ptr == '\\' || *ptr == '/') {
			ptr++;
			break;
		}
		if (ptr == filename) break;
	}
	if (ptr != filename) memmove(filename,ptr,(int)strlen(ptr)+1);
}

/* API Functions */
int pslib_init(void) {
	/* Load psapi.dll */
	if (!hInstLib) {
		hInstLib = LoadLibrary("PSAPI.DLL");
		if(hInstLib == NULL) {
			pslib_set_error();
			return PSLIB_LIBRARY;
		}
	}
	/* Get fn pointers */
	if (!lpfEnumProcesses) {
		lpfEnumProcesses = (BOOL(WINAPI *)(DWORD *,DWORD,DWORD*)) GetProcAddress(hInstLib,"EnumProcesses");
		if (lpfEnumProcesses == NULL) {
			pslib_set_error();
			return PSLIB_ENUM_PROCESSES;
		}
	}
	if (!lpfEnumProcessModules) {
		lpfEnumProcessModules = (BOOL(WINAPI *)(HANDLE, HMODULE *, DWORD, LPDWORD)) GetProcAddress(hInstLib,"EnumProcessModules");
		if (lpfEnumProcessModules == NULL) {
			pslib_set_error();
			return PSLIB_ENUM_PROCESS_MODULES;
		}
	}
	if (!lpfGetModuleFileNameEx) {
		lpfGetModuleFileNameEx =(DWORD (WINAPI *)(HANDLE, HMODULE, LPTSTR, DWORD )) GetProcAddress(hInstLib,"GetModuleFileNameExA");
		if (lpfGetModuleFileNameEx == NULL) {
			pslib_set_error();
			return PSLIB_GET_MODULE_FILENAME;
		}
	}
	initialized = TRUE;
	return PSLIB_SUCCESS;
}

void pslib_free(void) {
	if (hInstLib != NULL) {
		initialized = FALSE;
		lpfGetModuleFileNameEx = NULL;
		lpfEnumProcessModules = NULL;
		lpfEnumProcesses = NULL;
		FreeLibrary(hInstLib);
		hInstLib = NULL;		
	}
}

int pslib_list(int (*callbk)(void *,char *,int), void *data) {
	DWORD pids[PSLIB_MAX_LIST];
	char filename[MAX_PATH];	
	int count,index;
	HANDLE hProcess;
	HMODULE hMod;
	DWORD total;

	if (!callbk) return PSLIB_SUCCESS;

	if (!lpfEnumProcesses) return PSLIB_ENUM_PROCESSES;
	if (!lpfEnumProcessModules) return PSLIB_ENUM_PROCESS_MODULES;
	if (!lpfGetModuleFileNameEx) return PSLIB_GET_MODULE_FILENAME;

	if(!lpfEnumProcesses(pids,sizeof(DWORD)*PSLIB_MAX_LIST,&total)) {
		pslib_set_error();
		return PSLIB_ENUM_FAILED;
	}
	count = total/sizeof(DWORD);

	for (index = 0; index < count; index++) {
		hProcess = OpenProcess(PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,FALSE,pids[index]);
		if( hProcess == NULL ) continue;

		if(!lpfEnumProcessModules(hProcess,&hMod,sizeof(hMod),&total)) {
			CloseHandle(hProcess);
			continue;
		}
		if (!lpfGetModuleFileNameEx(hProcess,hMod,filename,MAX_PATH)) {
			CloseHandle(hProcess);
			continue;
		}
		CloseHandle(hProcess);
		strip_path(filename);
		if (callbk(data,filename,pids[index])==FALSE) break;
	}

	return PSLIB_SUCCESS;
}

int pslib_priority(int pid, DWORD priority) {
	HANDLE hProcess;

	hProcess = OpenProcess(PROCESS_SET_INFORMATION,FALSE,pid);
	if(hProcess == NULL) {
		pslib_set_error();
		return PSLIB_PRIORITY_PERMISSION;
	}
	if (SetPriorityClass(hProcess,priority) == 0) {
		pslib_set_error();
		return PSLIB_PRIORITY_FAILED;
	}
	return PSLIB_SUCCESS;
}

int pslib_terminate(int pid) {
	HANDLE hProcess;
	
	hProcess = OpenProcess(PROCESS_TERMINATE,FALSE,pid);
	if(hProcess == NULL) {
		pslib_set_error();
		return PSLIB_TERMINATE_PERMISSION;
	}
	if (TerminateProcess(hProcess,0) == 0) {
		pslib_set_error();
		return PSLIB_TERMINATE_FAILED;
	}
	return PSLIB_SUCCESS;
}

int pslib_find(char *filename, int *pid) {	
	int retcode;
	fData data;	

	if (!filename) return PSLIB_FIND_NOTARGET;
	
	memset(&data,0,sizeof(fData));
	data.filename = filename;
	
	retcode = pslib_list((int(*)(void *,char *,int))find_callbk,(void *)&data);
	if (retcode != PSLIB_SUCCESS) return retcode;

	*pid = data.pid;

	return PSLIB_SUCCESS;
}

char *pslib_error(int errcode, char *result) {
	if (errcode == PSLIB_LIBRARY)
		sprintf(result,"Failed to load psapi.dll, %s",system_error_string);
	else if (errcode == PSLIB_ENUM_PROCESSES)
		sprintf(result,"Failed to get address of EnumProcesses function, %s",system_error_string);
	else if (errcode == PSLIB_ENUM_PROCESS_MODULES)
		sprintf(result,"Failed to get address of EnumProcessModules function, %s",system_error_string);
	else if (errcode == PSLIB_GET_MODULE_FILENAME)
		sprintf(result,"Failed to get address of GetModuleFileNameExA function, %s",system_error_string);
	else if (errcode == PSLIB_ENUM_FAILED)
		sprintf(result,"Failed to enumerate processes, %s",system_error_string);
	else if (errcode == PSLIB_PRIORITY_FAILED)
		sprintf(result,"Failed to set process priority, %s",system_error_string);
	else if (errcode == PSLIB_TERMINATE_PERMISSION)
		sprintf(result,"Failed to get handle on process for terminate, %s",system_error_string);
	else if (errcode == PSLIB_TERMINATE_FAILED)
		sprintf(result,"Failed to terminate process, %s",system_error_string);
	return result;
}

int pslib_lower_priority(char *process, char *reason, int bfsz)
{
	char bf[PSLIB_ERROR_BFSZ];
	lPdata pdata;
	int retcode;
	
	memset(&pdata,0,sizeof(lPdata));
	pdata.process = process;

	retcode = pslib_list((int(*)(void *,char *,int))pslib_cb_lower,(void *)&pdata);
	if (retcode != PSLIB_SUCCESS) {
		lprintf(reason,bfsz,"pslib_list: %s\n",pslib_error(retcode,bf));
		return FALSE;
	}

	if (pdata.retcode != PSLIB_SUCCESS) {
		lprintf(reason,bfsz,"callback: %s\n",pslib_error(pdata.retcode,bf));
		return FALSE;
	}

	return TRUE;
}
