// By Cypher. Website: http://www.cypherjb.com/. Email: cypher.jb@gmail.com

// Custom RtlSetProcessIsCritical implementation, extended to handle 
// arbitrary processes.
NTSTATUS RtlSetProcessIsCriticalEx(HANDLE Process, BOOLEAN NewValue, 
	PBOOLEAN OldValue, BOOLEAN NeedBreaks)
{
  // Clear old 'BreakOnTermination' value if specified
  if (OldValue)
  {
    *OldValue = 0;
  }

  // If requested check whether system critical breaks are enabled
  PPEB pPeb = RtlGetCurrentPeb();
  ULONG FLG_ENABLE_SYSTEM_CRIT_BREAKS = 0x100000;
  if (NeedBreaks && Process == GetCurrentProcess() && 
    (pPeb->NtGlobalFlag & FLG_ENABLE_SYSTEM_CRIT_BREAKS) == 0)
  {
    return STATUS_UNSUCCESSFUL;
  }
  // Checks for system critical breaks disabled
  else
  {
    // Get old 'BreakOnTermination' value if specified
    if (OldValue)
    {
      ULONG OldValueTmp = 0;
      NtQueryInformationProcess(Process, ProcessBreakOnTermination, 
        &OldValueTmp, sizeof(OldValueTmp), 0);
      *OldValue = static_cast<BOOLEAN>(OldValueTmp);
    }

    // Set new 'BreakOnTermination' value
    ULONG NewValueTmp = NewValue;
    return NtSetInformationProcess(Process, ProcessBreakOnTermination, 
      &NewValueTmp, sizeof(NewValueTmp));
  }
}

// Gets the SeDebugPrivilege
void GetSeDebugPrivilege()
{
  // Open current process token with adjust rights
  HANDLE TempToken = 0;
  BOOL const RetVal = OpenProcessToken(GetCurrentProcess(), 
    TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &TempToken);
  if (!RetVal) 
  {
    throw std::runtime_error("Could not open process token.");
  }
  EnsureCloseHandle const Token(TempToken);

  // Get the LUID for SE_DEBUG_NAME 
  LUID Luid = { NULL }; // Locally unique identifier
  if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &Luid)) 
  {
    throw std::runtime_error("Could not look up privilege value for "
    	"SeDebugName.");
  }
  if (Luid.LowPart == NULL && Luid.HighPart == NULL) 
  {
    throw std::runtime_error("Could not get LUID for SeDebugName.");
  }

  // Process privileges
  TOKEN_PRIVILEGES Privileges = { NULL };
  // Set the privileges we need
  Privileges.PrivilegeCount = 1;
  Privileges.Privileges[0].Luid = Luid;
  Privileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

  // Apply the adjusted privileges
  if (!AdjustTokenPrivileges(Token, FALSE, &Privileges, sizeof (Privileges), 
    NULL, NULL)) 
  {
    throw std::runtime_error("Could not adjust token privileges.");
  }
}