/******************************************************************************
Module:  ProcessInfo.cpp
Notices: Copyright (c) 2008 Jeffrey Richter & Christophe Nasarre
******************************************************************************/


#include "..\CommonFiles\CmnHdr.h"     /* See Appendix A. */
#include "..\CommonFiles\Toolhelp.h"
#include <windowsx.h>
#include <stdarg.h>
#include <stdio.h>
#include "Resource.h"

#include <winternl.h>   // for Windows internal declarations.
#include <aclapi.h>     // for ACL management.
#include <shlwapi.h>    // for StrFormatKBSize.
#include <shlobj.h>     // for IsUserAnAdmin.
#include <AclApi.h>     // for ACL/ACE functions.

#include <tchar.h>
#include <StrSafe.h> 


#pragma comment (lib,"shlwapi.lib")
#pragma comment (lib,"shell32.lib")



// static variables
TOKEN_ELEVATION_TYPE s_elevationType = TokenElevationTypeDefault;
BOOL                 s_bIsAdmin = FALSE;
const int			 s_cchAddress = sizeof(PVOID) * 2;


///////////////////////////////////////////////////////////////////////////////


// Add a string to an edit control
void AddText(HWND hwnd, PCTSTR pszFormat, ...) {

   va_list argList;
   va_start(argList, pszFormat);

   TCHAR sz[20 * 1024];
   Edit_GetText(hwnd, sz, _countof(sz));
   _vstprintf_s(_tcschr(sz, TEXT('\0')), _countof(sz) - _tcslen(sz), 
      pszFormat, argList);
   Edit_SetText(hwnd, sz);
   va_end(argList);
}


///////////////////////////////////////////////////////////////////////////////


BOOL GetProcessIntegrityLevel(HANDLE hProcess, PDWORD pIntegrityLevel, 
   PDWORD pPolicy, PDWORD pResourceIntegrityLevel, PDWORD pResourcePolicy) {
   
   HANDLE hToken = NULL;
   if (!OpenProcessToken(hProcess, TOKEN_READ, &hToken)) {
      return(FALSE);
   }

   BOOL bReturn = FALSE;
   
   // First, compute the size of the buffer to get the Integrity level
   DWORD dwNeededSize = 0;
   if (!GetTokenInformation(
      hToken, TokenIntegrityLevel, NULL, 0, &dwNeededSize)) {

      PTOKEN_MANDATORY_LABEL pTokenInfo = NULL;
      if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
         // Second, allocate a memory block with the the required size 
         pTokenInfo = (PTOKEN_MANDATORY_LABEL)LocalAlloc(0, dwNeededSize);
         if (pTokenInfo != NULL) {
            // And finally, ask for the integrity level
            if (GetTokenInformation(hToken, TokenIntegrityLevel, pTokenInfo, 
               dwNeededSize, &dwNeededSize)) {

               *pIntegrityLevel = 
                  *GetSidSubAuthority(
                     pTokenInfo->Label.Sid, 
                     (*GetSidSubAuthorityCount(pTokenInfo->Label.Sid)-1)
                     );
               bReturn = TRUE;
            }
          
            // Don't forget to free the memory
            LocalFree(pTokenInfo);
         }
      }
   }

   // Try to get the policy if the integrity level was available
   if (bReturn) {
      *pPolicy = TOKEN_MANDATORY_POLICY_OFF;
      dwNeededSize = sizeof(DWORD);
      GetTokenInformation(hToken, TokenMandatoryPolicy, pPolicy, 
         dwNeededSize, &dwNeededSize);
   }
   
   // Look for the resource policy
   *pResourceIntegrityLevel = 0; // 0 means none explicitely set
   *pResourcePolicy = 0;
   
   PACL pSACL = NULL;
   PSECURITY_DESCRIPTOR pSD = NULL;
   DWORD dwResult = ERROR_SUCCESS;
   
   // Look for the no-read-up/no-write-up policy in the SACL
   if (hToken != NULL) {
      dwResult = 
         GetSecurityInfo(
            hProcess, SE_KERNEL_OBJECT,
            LABEL_SECURITY_INFORMATION,
            NULL, NULL, NULL, 
            &pSACL, &pSD
          );
      if (dwResult == ERROR_SUCCESS) {
         if (pSACL != NULL) {
            SYSTEM_MANDATORY_LABEL_ACE* pACE = NULL;
            if ((pSACL->AceCount > 0) && (GetAce(pSACL, 0, (PVOID*)&pACE))) {
               if (pACE != NULL) {
                  SID* pSID = (SID*)(&pACE->SidStart);
                  *pResourceIntegrityLevel = pSID->SubAuthority[0];
                  *pResourcePolicy = pACE->Mask;
               }
            }
         }
      }
      
      // Cleanup memory allocated on our behalf
      if (pSD != NULL) LocalFree(pSD); 
   }


   // Don't forget to close the token handle.
   CloseHandle(hToken);

   return(bReturn);   
}

BOOL GetProcessIntegrityLevel(DWORD PID, PDWORD pIntegrityLevel, 
   PDWORD pPolicy, PDWORD pResourceIntegrityLevel, PDWORD pResourcePolicy) {
   
   // Sanity checks
   if ((PID <= 0) || (pIntegrityLevel == NULL))
      return(FALSE);

   // Check if we can get information for this process
   HANDLE hProcess = OpenProcess(
      READ_CONTROL | PROCESS_QUERY_INFORMATION, 
      FALSE, PID);

   if (hProcess == NULL)
      return(FALSE);

   BOOL bReturn = GetProcessIntegrityLevel(hProcess, pIntegrityLevel, 
      pPolicy, pResourceIntegrityLevel, pResourcePolicy);

   // Don't forget to release the process handle
   CloseHandle(hProcess);

   return(bReturn);
}



VOID Dlg_PopulateProcessList(HWND hwnd) {

   HWND hwndList = GetDlgItem(hwnd, IDC_PROCESSMODULELIST);
   SetWindowRedraw(hwndList, FALSE);
   ComboBox_ResetContent(hwndList);

   CToolhelp thProcesses(TH32CS_SNAPPROCESS);
   PROCESSENTRY32 pe = { sizeof(pe) };
   BOOL fOk = thProcesses.ProcessFirst(&pe);
   for (; fOk; fOk = thProcesses.ProcessNext(&pe)) {
      TCHAR sz[1024];

      // Place the process name (without its path) & ID in the list
      PCTSTR pszExeFile = _tcsrchr(pe.szExeFile, TEXT('\\'));
      if (pszExeFile == NULL) {
         pszExeFile = pe.szExeFile;
      } else {
         pszExeFile++; // Skip over the slash
      }

      // Append the code/resource integrity level and policy
      DWORD dwCodeIntegrityLevel = 0;
      DWORD dwCodePolicy = TOKEN_MANDATORY_POLICY_OFF;
      DWORD dwResourcePolicy = 0;
      DWORD dwResourceIntegrityLevel = 0;
      TCHAR szCodeDetails[256];
      szCodeDetails[0] = TEXT('\0');
      TCHAR szResourceDetails[256];
      szResourceDetails[0] = TEXT('\0');
      if (GetProcessIntegrityLevel(pe.th32ProcessID, &dwCodeIntegrityLevel, 
         &dwCodePolicy, &dwResourceIntegrityLevel, &dwResourcePolicy)) {
         switch (dwCodeIntegrityLevel) {
            case SECURITY_MANDATORY_LOW_RID:
               _tcscpy_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT("- Low "));
               break;

            case SECURITY_MANDATORY_MEDIUM_RID:
               _tcscpy_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT("- Medium "));
               break;

            case SECURITY_MANDATORY_HIGH_RID:
               _tcscpy_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT("- High "));
               break;

            case SECURITY_MANDATORY_SYSTEM_RID:
               _tcscpy_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT("- System "));
               break;

            default:
               _tcscpy_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT("- ??? "));
         }

         if (dwCodePolicy == TOKEN_MANDATORY_POLICY_OFF) { // = 0
            _tcscat_s(szCodeDetails, 
               _countof(szCodeDetails), TEXT(" + no policy"));
         } else {
            if ((dwCodePolicy & TOKEN_MANDATORY_POLICY_VALID_MASK) == 0) {
               _tcscat_s(szCodeDetails, _countof(szCodeDetails), 
                  TEXT(" + ???"));
            } else {
               if ((dwCodePolicy & TOKEN_MANDATORY_POLICY_NO_WRITE_UP)
                  == TOKEN_MANDATORY_POLICY_NO_WRITE_UP) { 
                  _tcscat_s(szCodeDetails, _countof(szCodeDetails), 
                     TEXT(" + no write-up"));
               }

               if ((dwCodePolicy & TOKEN_MANDATORY_POLICY_NEW_PROCESS_MIN)
                  == TOKEN_MANDATORY_POLICY_NEW_PROCESS_MIN) { 
                  _tcscat_s(szCodeDetails, _countof(szCodeDetails), 
                     TEXT(" + new process min"));
               }
            }
         }

         switch (dwResourceIntegrityLevel) {
            case SECURITY_MANDATORY_LOW_RID:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("Low"));
               break;

            case SECURITY_MANDATORY_MEDIUM_RID:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("Medium"));
               break;

            case SECURITY_MANDATORY_HIGH_RID:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("High"));
               break;

            case SECURITY_MANDATORY_SYSTEM_RID:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("System"));
               break;

            case 0:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("Not set"));
               break;

            default:
               _tcscpy_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT("???"));
          }


         if (dwResourcePolicy == 0) { // = 0
            _tcscat_s(szResourceDetails, 
               _countof(szResourceDetails), TEXT(" + 0 policy"));
         } else {
            if ((dwResourcePolicy & TOKEN_MANDATORY_POLICY_VALID_MASK) == 0) {
               _tcscat_s(szResourceDetails, 
                  _countof(szResourceDetails), TEXT(" + ???"));
            } else {
               if ((dwResourcePolicy & SYSTEM_MANDATORY_LABEL_NO_WRITE_UP)
                  == SYSTEM_MANDATORY_LABEL_NO_WRITE_UP) { 
                  _tcscat_s(szResourceDetails, 
                     _countof(szResourceDetails), 
                     TEXT(" + no write-up"));
               }

               if ((dwResourcePolicy & SYSTEM_MANDATORY_LABEL_NO_READ_UP)
                  == SYSTEM_MANDATORY_LABEL_NO_READ_UP) { 
                  _tcscat_s(szResourceDetails, 
                     _countof(szResourceDetails), 
                     TEXT(" + no read-up"));
               }
               if ((dwResourcePolicy & SYSTEM_MANDATORY_LABEL_NO_EXECUTE_UP)
                  == SYSTEM_MANDATORY_LABEL_NO_EXECUTE_UP) { 
                  _tcscat_s(szResourceDetails, 
                     _countof(szResourceDetails), 
                     TEXT(" + no execute-up"));
               }
            }
         }
      }

      StringCchPrintf(sz, _countof(sz), TEXT("%s     (0x%08X)  %s    [%s]"), 
         pszExeFile, pe.th32ProcessID, szCodeDetails, szResourceDetails);
      int n = ComboBox_AddString(hwndList, sz);

      // Associate the process ID with the added item
      ComboBox_SetItemData(hwndList, n, pe.th32ProcessID);
   }

   ComboBox_SetCurSel(hwndList, 0);  // Select the first entry

   // Simulate the user selecting this first item so that the
   // results pane shows something interesting
   FORWARD_WM_COMMAND(hwnd, IDC_PROCESSMODULELIST, 
      hwndList, CBN_SELCHANGE, SendMessage);

   SetWindowRedraw(hwndList, TRUE);
   InvalidateRect(hwndList, NULL, FALSE);
}


///////////////////////////////////////////////////////////////////////////////


VOID Dlg_PopulateModuleList(HWND hwnd) {

   HWND hwndModuleHelp = GetDlgItem(hwnd, IDC_MODULEHELP);
   ListBox_ResetContent(hwndModuleHelp);

   CToolhelp thProcesses(TH32CS_SNAPPROCESS);
   PROCESSENTRY32 pe = { sizeof(pe) };
   BOOL fOk = thProcesses.ProcessFirst(&pe);
   for (; fOk; fOk = thProcesses.ProcessNext(&pe)) {

      CToolhelp thModules(TH32CS_SNAPMODULE, pe.th32ProcessID);
      MODULEENTRY32 me = { sizeof(me) };
      BOOL fOk = thModules.ModuleFirst(&me);
      for (; fOk; fOk = thModules.ModuleNext(&me)) {
        int n = ListBox_FindStringExact(hwndModuleHelp, -1, me.szExePath);
         if (n == LB_ERR) {
            // This module hasn't been added before
            ListBox_AddString(hwndModuleHelp, me.szExePath);
         }
      }
   }

   HWND hwndList = GetDlgItem(hwnd, IDC_PROCESSMODULELIST);
   SetWindowRedraw(hwndList, FALSE);
   ComboBox_ResetContent(hwndList);
   int nNumModules = ListBox_GetCount(hwndModuleHelp);
   for (int i = 0; i < nNumModules; i++) {
      TCHAR sz[1024];
      ListBox_GetText(hwndModuleHelp, i, sz);
      // Place module name (without its path) in the list
      int nIndex = ComboBox_AddString(hwndList, _tcsrchr(sz, TEXT('\\')) + 1);
      // Associate the index of the full path with the added item
      ComboBox_SetItemData(hwndList, nIndex, i);
   }

   ComboBox_SetCurSel(hwndList, 0);  // Select the first entry

   // Simulate the user selecting this first item so that the
   // results pane shows something interesting
   FORWARD_WM_COMMAND(hwnd, IDC_PROCESSMODULELIST, 
      hwndList, CBN_SELCHANGE, SendMessage);

   SetWindowRedraw(hwndList, TRUE);
   InvalidateRect(hwndList, NULL, FALSE);
}