// PSLProcesses.cpp : Implementation of CPSLProcesses

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

ULONG CPSLProcesses::m_BufferSize = 0x8000; // 32KB by default;

CPSLProcesses::IsWow64ProcessProc CPSLProcesses::m_IsWow64Process = NULL;
HMODULE CPSLProcesses::m_hKernelModule = NULL;

CPSLProcesses::CPSLProcesses()
{
   TCHAR buffer[UNLEN + 1];
   DWORD dwSize = UNLEN + 1;
   ::GetUserName(buffer, &dwSize);
   m_sUserName = buffer;

   m_bDirty = true;
   m_EnumType = enumForCurrentUser;

   m_IsWow64Process = NULL;
   m_hKernelModule = ::LoadLibrary(_T("kernel32.dll"));
   if(m_hKernelModule)
      m_IsWow64Process = (IsWow64ProcessProc)::GetProcAddress(m_hKernelModule, "IsWow64Process");
}

CPSLProcesses::~CPSLProcesses()
{
   if(m_hKernelModule)
   {
      ::FreeLibrary(m_hKernelModule);
      m_hKernelModule = NULL;
      m_IsWow64Process = NULL;
   }
}

bool CPSLProcesses::Is64BitProcess(HANDLE hProcess)
{
   if(m_IsWow64Process)
   {
      BOOL bIs64Bit = FALSE;
      if(m_IsWow64Process(hProcess, &bIs64Bit))
         return bIs64Bit?false:true;
   }
   return true;
}

HRESULT CPSLProcesses::OnIndexOutOfRange()
{
   return MakeException(EX_INDEXOUTOFRANGE);
}

HRESULT CPSLProcesses::OnInterfaceAccess()
{
   InternalUpdate(false);
   return GetExitCode();
}

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

void CPSLProcesses::FinalRelease()
{
}

/*
Interesting: GetWindowModuleFileName
When/If implementing Windows property
for namespace Software, this could be an asset;
*/

void CPSLProcesses::InternalUpdate(bool bForceUpdate)
{
   CCritSecLock cs(m_csCollection);

   if(!bForceUpdate && !m_bDirty)
      return;

   m_bDirty = false;

   m_coll.clear(); // Releases interfaces and clears the collection;
               // NOTE: .NET clients catch up interfaces, so takes
               // about 30 seconds before they are garbaged by .NET;

   HANDLE hHeap = ::GetProcessHeap();
    NTSTATUS Status;
    PVOID pBuffer = NULL;

   LPCTSTR sUserName = NULL;
   if(m_EnumType == enumForCurrentUser)
      sUserName = m_sUserName;

    do
    {
      pBuffer = ::HeapAlloc(hHeap, 0, m_BufferSize);
      if(pBuffer == NULL)
      {
         SetException(EX_LOWMEMORY);
         return; // Not enough memory;
      }

        Status = g_SIA.ZwQuerySystemInformation(SystemProcessesAndThreadsInformation, pBuffer, m_BufferSize);
        if(Status == STATUS_INFO_LENGTH_MISMATCH) // Not enough memory was allocated;
        {
         ::HeapFree(hHeap, 0, pBuffer);
            m_BufferSize *= 2; // Let's try and increase it by 2;
        }
      else
         if(!NT_SUCCESS(Status)) // If still failed;
         {
            ::HeapFree(hHeap, 0, pBuffer);
            SetException(EX_GENERALFAILURE);
            return; // There was enough memory, but the call still failed with Status;
         }
    }
    while(Status == STATUS_INFO_LENGTH_MISMATCH);
    PSYSTEM_PROCESS_INFORMATION pProcesses = (PSYSTEM_PROCESS_INFORMATION)pBuffer;
   CComObject<CPSLProcess> * pProcess = NULL;
    do
    {
      if(!pProcess)
         CComObject<CPSLProcess>::CreateInstance(&pProcess);
      if(pProcess && pProcess->Initialize(pProcesses, sUserName))
      {
         m_coll.push_back(CComPtr<IPSLProcess>(pProcess));
         pProcess = NULL;
      }
      if(pProcesses->NextEntryDelta == 0) // If the last process;
         break;
        // find the address of the next process structure
        pProcesses = (PSYSTEM_PROCESS_INFORMATION)(((LPBYTE)pProcesses) + pProcesses->NextEntryDelta);
    }
   while(1);
   if(pProcess)
      pProcess->Release();
   ::HeapFree(hHeap, 0, pBuffer);
}

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

STDMETHODIMP CPSLProcesses::Update()
{
   PSL_BEGIN

   InternalUpdate(true);

   PSL_END
}

STDMETHODIMP CPSLProcesses::Find(long ProcessID, IPSLProcess ** ppValue)
{
   PSL_BEGIN

   *ppValue = NULL;
   InternalUpdate(false);
   CCritSecLock cs(m_csCollection);
   for(ProcessesType::iterator i = m_coll.begin();i != m_coll.end();i ++)
   {
      long ID = 0;
      if(i->m_T->get_ProcessID(&ID) == S_OK && ID == ProcessID)
      {
         IPSLProcess * pProcess = i->m_T;
         pProcess->AddRef();
         *ppValue = pProcess;
         break;
      }
   }

   PSL_END
}

STDMETHODIMP CPSLProcesses::get_EnumUserType(PSLEnumUserType * pValue)
{
   PSL_BEGIN

   CCritSecLock cs(m_csCollection);
   *pValue = m_EnumType;

   PSL_END
}

STDMETHODIMP CPSLProcesses::put_EnumUserType(PSLEnumUserType newValue)
{
   PSL_BEGIN

   CCritSecLock cs(m_csCollection);
   if(m_EnumType != newValue)
   {
      m_EnumType = newValue;
      m_bDirty = true; // Collection needs to be updated;
   }

   PSL_END
}