#include "precomp.h"


static KSPIN_LOCK s_lock;

//#define PID_MAP

static ULONG s_allowPids[1024];
static ULONG s_denyPids[1024];
//static ULONG s_protectedPids[512];

static ULONG s_whitePids[1024];
static ULONG s_greenPids[1024];

static ULONG s_protectPids[1024];

static void SetPidBitInMap(ULONG* m, ULONG pid)
{
	ULONG rpid = (ULONG)(pid & 0xffffff) >> 2;
	ULONG i = rpid >> 5;
	ULONG r = rpid & 0x1f;
	ULONG mask = 1 << (31 - r);

	m[i] |= mask;
}

static ULONG GetPidBitInMap(ULONG* m, ULONG pid)
{
	ULONG p = (ULONG)(pid & 0xffffff) >> 2;
	ULONG i = p >> 5;
	ULONG r = p & 0x1f;
	ULONG q = m[i] << r;

	return q & 0x80000000;
}

static void ClearPidBitInMap(ULONG* m, ULONG pid)
{
	ULONG rpid = (ULONG)(pid & 0xffffff) >> 2;
	ULONG i = rpid >> 5;
	ULONG r = rpid & 0x1f;
	ULONG mask = 1 << (31 - r);

	m[i] &= ~mask;
}


void PgInitialize()
{
	KeInitializeSpinLock(&s_lock);
	RtlZeroMemory(s_allowPids, sizeof(s_allowPids));
	RtlZeroMemory(s_whitePids, sizeof(s_whitePids));
	RtlZeroMemory(s_greenPids, sizeof(s_greenPids));
	RtlZeroMemory(s_denyPids, sizeof(s_denyPids));
	RtlZeroMemory(s_protectPids, sizeof(s_denyPids));

}

void PgCleanup()
{
}

void PgReset()
{
	KIRQL irql;
	KeAcquireSpinLock(&s_lock, &irql);
	RtlZeroMemory(s_allowPids, sizeof(s_allowPids));
	RtlZeroMemory(s_whitePids, sizeof(s_whitePids));
	RtlZeroMemory(s_greenPids, sizeof(s_greenPids));
	RtlZeroMemory(s_denyPids, sizeof(s_denyPids));
	RtlZeroMemory(s_protectPids, sizeof(s_denyPids));

	KeReleaseSpinLock(&s_lock, irql);
}

void PgResetState(ULONG state)
{
	KIRQL irql;
	KeAcquireSpinLock(&s_lock, &irql);
	if (state == PG_PID_ALLOW)
		RtlZeroMemory(s_allowPids, sizeof(s_allowPids));
	if (state == PG_PID_WHITE)
		RtlZeroMemory(s_whitePids, sizeof(s_whitePids));
	if (state == PG_PID_GREEN)
		RtlZeroMemory(s_greenPids, sizeof(s_greenPids));
	if (state == PG_PID_BLACK)
		RtlZeroMemory(s_denyPids, sizeof(s_denyPids));
	if (state == PG_PID_PROTECT)
		RtlZeroMemory(s_protectPids, sizeof(s_denyPids));

	KeReleaseSpinLock(&s_lock, irql);
}

void PgAddPid(ULONG pid, ULONG state)
{
	KIRQL irql;

	//KLogEx("%d %d\n", pid, state);

	KeAcquireSpinLock(&s_lock, &irql);
	if (state == PG_PID_ALLOW)
		SetPidBitInMap(s_allowPids, pid);
	if (state == PG_PID_WHITE)
		SetPidBitInMap(s_whitePids, pid);
	if (state == PG_PID_GREEN)
		SetPidBitInMap(s_greenPids, pid);
	if (state == PG_PID_BLACK)
		SetPidBitInMap(s_denyPids, pid);
	if( state == PG_PID_PROTECT)
		SetPidBitInMap(s_protectPids, pid);

	KeReleaseSpinLock(&s_lock, irql);

	KLogEx(DEBUG_TRACE_INFO, "%d %d COMPLETE\n", pid, state);
}


ULONG PgRemovePid(ULONG pid)
{
	KIRQL irql;
	ULONG state = PgGetPidState(pid);
	KLogEx(DEBUG_TRACE_INFO, "PgRemove %d %d\n", pid, state);

	KeAcquireSpinLock(&s_lock, &irql);
	ClearPidBitInMap(s_allowPids, pid);
	ClearPidBitInMap(s_whitePids, pid);
	ClearPidBitInMap(s_greenPids, pid);
	ClearPidBitInMap(s_denyPids, pid);
	ClearPidBitInMap(s_protectPids, pid);

	KeReleaseSpinLock(&s_lock, irql);
	return state;
}

ULONG PgRemovePidState(ULONG pid, ULONG state)
{
	KIRQL irql;
	KLogEx(DEBUG_TRACE_INFO, "%d %d\n", pid, state);

	KeAcquireSpinLock(&s_lock, &irql);

	if (state == PG_PID_ALLOW)
		ClearPidBitInMap(s_allowPids, pid);
	if (state == PG_PID_WHITE)
		ClearPidBitInMap(s_whitePids, pid);
	if (state == PG_PID_GREEN)
		ClearPidBitInMap(s_greenPids, pid);
	if (state == PG_PID_BLACK)
		ClearPidBitInMap(s_denyPids, pid);
	if (state == PG_PID_PROTECT)
		ClearPidBitInMap(s_protectPids, pid);

	KeReleaseSpinLock(&s_lock, irql);
	return state;
}

ULONG PgGetProtectPidState(ULONG pid)
{
	KIRQL irql;
	ULONG aState;

	KeAcquireSpinLock(&s_lock, &irql);
	aState = GetPidBitInMap(s_protectPids, pid);
	KeReleaseSpinLock(&s_lock, irql);
	if (aState)
		return PG_PID_PROTECT;
	
	return PG_PID_UNDEFINED;

}
ULONG PgGetPidState(ULONG pid)
{
	KIRQL irql;
	ULONG aState, wState, gstate, bstate;

	KeAcquireSpinLock(&s_lock, &irql);
	aState = GetPidBitInMap(s_allowPids, pid);
	wState = GetPidBitInMap(s_whitePids, pid);
	gstate = GetPidBitInMap(s_greenPids, pid);
	bstate = GetPidBitInMap(s_denyPids, pid);
	
	KeReleaseSpinLock(&s_lock, irql);
	if (aState)
		return PG_PID_ALLOW;
	if (wState)
		return PG_PID_WHITE;
	if (gstate)
		return PG_PID_GREEN;
	if (bstate)
		return PG_PID_BLACK;

	return PG_PID_UNDEFINED;
}