// PSLProcess.cpp : Implementation of CPSLProcess

#include "stdafx.h"
#include "PSLProcess.h"
#include "PSLProcesses.h"
#include "Psapi.h"

CPSLProcess::CPSLProcess()
{
   Reset();
}

void CPSLProcess::Reset()
{
   m_ProcessID = 0;
   m_ThreadCount = 0;
   m_HandleCount = 0;
   m_GDICount = 0;
   m_USERCount = 0;
   m_Priority = ppUnknown;
   m_AffinityMask = 0;
   m_bIs64Bit = false;

   m_sFileName = _T("");
   m_sFilePath = _T("");
   m_sCommandLine = _T("");
   m_sUserName = _T("");
   m_sDomainName = _T("");
}

HRESULT CPSLProcess::FinalConstruct()
{
   return S_OK;
}

void CPSLProcess::FinalRelease()
{
}

_bstr_t CPSLProcess::FilePathFromProcessHandle(HANDLE hProcess)
{
   // Works on XP/Vista only, and retrieves logical path
   // instead of file path:
   //
   // tstring t;
   // t.resize(MAX_FILE_PATH);
   // GetProcessImageFileName(hProcess, (LPTSTR)t.c_str(), MAX_FILE_PATH);
   // return t.c_str();

   _bstr_t sResult(_T(""));
   tstring sPath;
   sPath.resize(MAX_FILE_PATH);
   if(::GetModuleFileNameEx(hProcess, NULL, (LPTSTR)sPath.c_str(), MAX_FILE_PATH))
      sResult = sPath.c_str();
   return sResult;
}

bool CPSLProcess::GetUserAndDomainFromSID(PSID pSid, tstring & sName, tstring & sDomain)
{
   DWORD dwNameLength = 0;
   DWORD dwDomainLength = 0;
   SID_NAME_USE NameUse;
   if(!::LookupAccountSid(NULL, pSid, NULL, &dwNameLength, NULL, &dwDomainLength, &NameUse) && ::GetLastError() == ERROR_INSUFFICIENT_BUFFER)
   {
      sName.resize(dwNameLength);
      sDomain.reserve(dwDomainLength);
      return ::LookupAccountSid(NULL, pSid, (LPTSTR)sName.c_str(), &dwNameLength, (LPTSTR)sDomain.c_str(), &dwDomainLength, &NameUse)?true:false;
   }
   return false;
}

bool CPSLProcess::Initialize(PSYSTEM_PROCESS_INFORMATION pProcess, LPCTSTR sUserName)
{
   Reset();

   m_ProcessID = (long)pProcess->ProcessId;
   m_ThreadCount = pProcess->ThreadCount;
   m_HandleCount = pProcess->HandleCount;

   if(pProcess->ProcessName.Buffer) // If not IDLE (System Idle Process) process;
      m_sFileName = pProcess->ProcessName.Buffer;

   HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION|PROCESS_VM_READ, FALSE, m_ProcessID);
   if(!hProcess)
      hProcess = ::OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, m_ProcessID);

   if(hProcess)
   {
      m_bIs64Bit = CPSLProcesses::Is64BitProcess(hProcess);
      m_sFilePath = FilePathFromProcessHandle(hProcess);
      m_GDICount = ::GetGuiResources(hProcess, GR_GDIOBJECTS);
      m_USERCount = ::GetGuiResources(hProcess, GR_USEROBJECTS);
      m_Priority = (PSLProcessPriority)::GetPriorityClass(hProcess);

      DWORD_PTR dwProcessAffinityMask = 0;
      DWORD_PTR dwSystemAffinityMask = 0;
      if(::GetProcessAffinityMask(hProcess, &dwProcessAffinityMask, &dwSystemAffinityMask))
         m_AffinityMask = dwProcessAffinityMask;

      HANDLE hToken = NULL;
      if(::OpenProcessToken(hProcess, TOKEN_READ, &hToken))
      {
         DWORD dwLength = 64;
         LPBYTE Buffer[64];
         if(::GetTokenInformation(hToken, TokenUser, Buffer, dwLength, &dwLength))
         {
            TOKEN_USER * pUser = (TOKEN_USER*)(LPVOID)Buffer;
            tstring sName, sDomain;
            if(GetUserAndDomainFromSID(pUser->User.Sid, sName, sDomain))
            {
               m_sUserName = sName.c_str();
               m_sDomainName = sDomain.c_str();
            }
         }
         ::CloseHandle(hToken);
      }
      ::CloseHandle(hProcess);
   }

   if(sUserName && ::_tcsicmp(sUserName, m_sUserName))
      return false; // Didn't pass the user filter;

   return true;
}

////////////////////////////////////////////////////////////////////////
// Interface Implementation;
////////////////////////////////////////////////////////////////////////

STDMETHODIMP CPSLProcess::get_FileName(BSTR * pValue)
{
   PSL_BEGIN

   *pValue = m_sFileName.copy();

   PSL_END
}

STDMETHODIMP CPSLProcess::get_FilePath(BSTR * pValue)
{
   PSL_BEGIN

   *pValue = m_sFilePath.copy();

   PSL_END
}

STDMETHODIMP CPSLProcess::get_CommandLine(BSTR * pValue)
{
   PSL_BEGIN

   SetException(EX_NOTIMPLEMENTED);
   //*pValue = m_sCommandLine.copy();

   PSL_END
}

STDMETHODIMP CPSLProcess::get_UserName(BSTR * pValue)
{
   PSL_BEGIN

   *pValue = m_sUserName.copy();

   PSL_END
}

STDMETHODIMP CPSLProcess::get_DomainName(BSTR * pValue)
{
   PSL_BEGIN

   *pValue = m_sDomainName.copy();

   PSL_END
}

STDMETHODIMP CPSLProcess::get_ProcessID(long * pValue)
{
   PSL_BEGIN

   *pValue = m_ProcessID;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_Is64Bit(VARIANT_BOOL * pValue)
{
   PSL_BEGIN

   *pValue = m_bIs64Bit?VARIANT_TRUE:VARIANT_FALSE;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_AffinityMask(VARIANT * pValue)
{
   PSL_BEGIN

   CPSLUtilities::SetVariantBig(pValue, 0);

   HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, m_ProcessID);
   if(hProcess)
   {
      DWORD_PTR ProcessAffinityMask = 0;
      DWORD_PTR SystemAffinityMask = 0;
      if(::GetProcessAffinityMask(hProcess, &ProcessAffinityMask, &SystemAffinityMask))
         m_AffinityMask = ProcessAffinityMask;
      CPSLUtilities::SetVariantBig(pValue, m_AffinityMask);
   }
   else
   {
      CPSLUtilities::SetVariantBig(pValue, m_AffinityMask);
      if(::GetLastError() == ERROR_ACCESS_DENIED)
         SetException(EX_NOACCESS);
   }

   PSL_END
}

STDMETHODIMP CPSLProcess::put_AffinityMask(VARIANT newValue)
{
   PSL_BEGIN

   _variant_t v = newValue;
   ABIG Mask = 0;

   try
   {
      Mask = v;
   }
   catch(...)
   {
   }
   if(Mask)
   {
      HANDLE hProcess = ::OpenProcess(PROCESS_SET_INFORMATION, FALSE, m_ProcessID);
      if(hProcess)
      {
         DWORD_PTR ProcessAffinityMask = Mask;
         ::SetProcessAffinityMask(hProcess, ProcessAffinityMask);
         ::CloseHandle(hProcess);
      }
      else
      {
         if(::GetLastError() == ERROR_ACCESS_DENIED)
            SetException(EX_NOACCESS);
      }
   }
   else
      SetException(EX_INVALIDPARAMETER);

   PSL_END
}

STDMETHODIMP CPSLProcess::get_ThreadCount(long * pValue)
{
   PSL_BEGIN

   *pValue = m_ThreadCount;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_HandleCount(long * pValue)
{
   PSL_BEGIN

   *pValue = m_HandleCount;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_GDICount(long * pValue)
{
   PSL_BEGIN

   *pValue = m_GDICount;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_USERCount(long * pValue)
{
   PSL_BEGIN

   *pValue = m_USERCount;

   PSL_END
}

STDMETHODIMP CPSLProcess::get_Priority(PSLProcessPriority * pValue)
{
   PSL_BEGIN

   *pValue = m_Priority;

   PSL_END
}

STDMETHODIMP CPSLProcess::Kill(VARIANT_BOOL * pValue)
{
   PSL_BEGIN

   *pValue = VARIANT_FALSE;
   HANDLE hProcess = ::OpenProcess(SYNCHRONIZE|PROCESS_TERMINATE, FALSE, m_ProcessID);
   if(hProcess)
   {
      if(::WaitForSingleObject(hProcess, 500) != WAIT_OBJECT_0)
      {
         *pValue = ::TerminateProcess(hProcess, 0)?VARIANT_TRUE:VARIANT_FALSE;
         if(!*pValue && ::GetLastError() == ERROR_ACCESS_DENIED)
            SetException(EX_NOACCESS);
      }
   }
   else
      if(::GetLastError() == ERROR_ACCESS_DENIED)
         SetException(EX_NOACCESS);

   PSL_END
}