/*
 * pmpctl.c
 *
 * Author: Andy Ying
 *   Date: April 20, 2007
 */

#include "pmpctl.h"

int main(int argc, char **argv)
{
	HANDLE hFile;
	
	char *cmd;
	unsigned long pid;

	if (argc != 3) {
		printf("pmpctl: invalid parameters\n");
		help();

		return -1;
	}

	cmd = *(argv + 1);
	pid = (unsigned int) atoi(*(argv + 2));

	printf("pmpctl: pid %u\n", pid);

	hFile = CreateFile("\\\\.\\PmpFun",
					   FILE_ALL_ACCESS,
					   FILE_SHARE_READ | FILE_SHARE_WRITE,
					   NULL,
					   OPEN_EXISTING,
					   FILE_FLAG_OVERLAPPED,
					   NULL);

	if (hFile == INVALID_HANDLE_VALUE) {
		printf("pmpctl: failed to open \\\\.\\PmpFun\n");
		printf("pmpctl: GetLastError() = %i\n", GetLastError());

		return -1;	
	}

	if (strncmp(cmd, "set", 3) == 0) {
		// Set PMP on a process
		if (!DeviceIoControl(hFile,
			     			 IOCTL_SET_PMP,
							 &pid,
							 sizeof(pid),
							 NULL,
							 0,
							 NULL,
							 NULL)) {
			printf("pmpctl: failed to set PMP on %u\n", pid);
			
			CloseHandle(hFile);

			return -1;
		}

		printf("pmpctl: PID %u is now protected!\n");
	} else if (strncmp(cmd, "reset", 5) == 0) {
		// Reset PMP on a process
		if (!DeviceIoControl(hFile,
			     			 IOCTL_RESET_PMP,
							 &pid,
							 sizeof(pid),
							 NULL,
							 0,
							 NULL,
							 NULL)) {
			printf("pmpctl: failed to reset PMP on %u\n", pid);
			
			CloseHandle(hFile);

			return -1;
		}

		printf("pmpctl: PID %u is now unprotected!\n");
	} else if (strncmp(cmd, "get", 3) == 0) {
		// Get PMP status
		unsigned int ret;
		
		if (!DeviceIoControl(hFile,
			     			 IOCTL_GET_PMP,
							 &pid,
							 sizeof(pid),
							 &ret,
							 sizeof(ret),
							 NULL,
							 NULL)) {
			printf("pmpctl: failed to get PMP status for %u\n", pid);
			
			CloseHandle(hFile);

			return -1;
		}

		printf("pmpctl: PID %u is has flags 0x%x\n", pid, ret);
	} else {
		printf("pmpctl: invalid command\n");
		help();
		
		CloseHandle(hFile);

		return -1;
	}
	
	CloseHandle(hFile);

	return 0;
}

void help() 
{
	printf("pmpctl [cmd] [pid]\n");
	printf("\t[cmd] - can be set, reset, get\n");
	printf("\t[pid] - PID of the process\n");
}