#include <windows.h>
#include <stdio.h>
#include "privs.h"
#include "error.h"

BOOL SetPrivilege(
	HANDLE hToken,          // token handle (NULL: current process)
    LPCTSTR Privilege,      // Privilege to enable/disable
    BOOL bEnablePrivilege   // TRUE to enable.  FALSE to disable
    )
{
    TOKEN_PRIVILEGES tp;
    LUID luid;
    TOKEN_PRIVILEGES tpPrevious;
    DWORD cbPrevious=sizeof(TOKEN_PRIVILEGES);
    BOOL CloseAtEnd=FALSE;

    //
    // Retrieve a handle of the access token
    //
    if (hToken==NULL) {
		if (!OpenProcessToken(GetCurrentProcess(),
				TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
				&hToken)) {
			ErrorHandler("OpenProcessToken in SetPrivilege");
			return FALSE;
		}
        CloseAtEnd=TRUE;
	}

    if (!LookupPrivilegeValue( NULL, Privilege, &luid )) {
		ErrorHandler("LookupPrivilegeValue in SetPrivilege");
		return FALSE;
	}

    //
    // first pass.  get current privilege setting
    //
    tp.PrivilegeCount           = 1;
    tp.Privileges[0].Luid       = luid;
    tp.Privileges[0].Attributes = 0;

    AdjustTokenPrivileges(
            hToken,
            FALSE,
            &tp,
            sizeof(TOKEN_PRIVILEGES),
            &tpPrevious,
            &cbPrevious
            );

    if (GetLastError() != ERROR_SUCCESS) {
		ErrorHandler("AdjustTokenPrivileges in SetPrivilege");
		return FALSE;
	}

    //
    // second pass.  set privilege based on previous setting
    //
    tpPrevious.PrivilegeCount       = 1;
    tpPrevious.Privileges[0].Luid   = luid;

    if(bEnablePrivilege) {
        tpPrevious.Privileges[0].Attributes |= (SE_PRIVILEGE_ENABLED);
    }
    else {
        tpPrevious.Privileges[0].Attributes ^= (SE_PRIVILEGE_ENABLED &
            tpPrevious.Privileges[0].Attributes);
    }

    AdjustTokenPrivileges(
            hToken,
            FALSE,
            &tpPrevious,
            cbPrevious,
            NULL,
            NULL
            );

    if (GetLastError() != ERROR_SUCCESS) {
		ErrorHandler("AdjustTokenPrivileges in SetPrivilege");
		return FALSE;
	}

    /* call CloseHandle or not ??? */
    if (CloseAtEnd) {
        CloseHandle(hToken);
    }
    
    return TRUE;
}

BOOL
EnableDebugPriv(
    VOID
    )
{
    return SetPrivilege(NULL,SE_DEBUG_NAME,TRUE);
}