#pragma region Includes
#include "SSAservice.h"
#include "ThreadPool.h"
#include "PathName.h"

#include <psapi.h>
#pragma comment(lib, "Psapi.lib")

#pragma endregion

///////////////////////////////////////////////////////////
// CSSAservice
///////////////////////////////////////////////////////////

BOOL CSSAservice::FindAgent(HANDLE& hFound, DWORD& pid) const
{
  DWORD* aProcesses = NULL;
  DWORD nItems = 0, nFound = 0, i = 0;
  
  hFound = 0;
  pid = 0;

  for (nItems = 256; ; nItems *= 2)
  {
    DWORD cbNeeded = 0;
    free(aProcesses);
    aProcesses = (DWORD*)calloc(nItems, sizeof(DWORD));
    if (!EnumProcesses(aProcesses, nItems*sizeof(DWORD), &cbNeeded))
    {
      free(aProcesses);
      return FALSE;
    }
    nFound = cbNeeded / sizeof(DWORD);
    if (nFound < nItems)
      break;
  }

  for (i = 0; i < nFound && hFound == NULL; i++) if (aProcesses[i])
  {
    HANDLE hProcess = ::OpenProcess( PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, aProcesses[i] );
    if (hProcess != NULL)
    {
      HMODULE hMod;
      DWORD cbNeeded;
      if (::EnumProcessModules( hProcess, &hMod, sizeof(hMod), &cbNeeded) )
      {
        TCHAR szProcessName[MAX_PATH] = { 0 };
        ::GetModuleBaseName( hProcess, hMod, szProcessName, sizeof(szProcessName)/sizeof(TCHAR) );
        if (wcsstr(szProcessName, _agent) != NULL)
        {
          hFound = ::OpenProcess( SYNCHRONIZE, FALSE, aProcesses[i] ); // Waitable process handle
          pid = aProcesses[i];
        }
      }
      // Release the handle to the process.
      CloseHandle( hProcess );
    }
  }

  free(aProcesses);
  return hFound != NULL;
}

BOOL CSSAservice::FindOrCreateAgent(HANDLE& hProcess, DWORD& pid) const
{
  if (!FindAgent(hProcess, pid))
  {
    STARTUPINFO si = { 0 };
    PROCESS_INFORMATION pi = { 0 };
    si.cb = sizeof(si); // si.wShowWindow = SW_HIDE;
  
    std::wstring str;
    str = L"Creating instance of "; str += _agent;
    WriteEventLogEntry(str.c_str());
    if (CreateProcess(_agent, 
                      NULL,  // lpCommandLine
                      NULL,  // lpProcessAttributes
                      NULL,  // lpThreadAttributes
                      FALSE, // no Handle inheritance
                      CREATE_NO_WINDOW | IDLE_PRIORITY_CLASS, // creation flags
                      NULL,  // Use parent's environment block
                      NULL,  // Use parent's starting directory 
                      &si, 
                      &pi))
    {
      hProcess = pi.hProcess;
      pid = pi.dwProcessId;
    }
    else
      WriteErrorLogEntry(TEXT("CreateProcess"));
  }
  else
  {
    std::wstring str;
    str = L"Attaching to instance of "; str += _agent;
    WriteEventLogEntry(str.c_str());
  }
  return hProcess != 0;
}

void CSSAservice::CheckDirectory() const
{
  char full[_MAX_PATH];
  ::GetModuleFileNameA(NULL, full, sizeof(full));
  CPathName pn = full;
  ::SetCurrentDirectoryA(pn.Path().c_str());
}

CSSAservice::CSSAservice(LPCTSTR pszServiceName, 
                         BOOL fCanStop, BOOL fCanShutdown, BOOL fCanPauseContinue)
: CServiceBase(pszServiceName, fCanStop, fCanShutdown, fCanPauseContinue), _agent(NULL)
{
    m_fStopping = FALSE;

    // Create a manual-reset event that is not signaled at first to indicate 
    // the stopped signal of the service.
    m_hStoppedEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (m_hStoppedEvent == NULL)
        throw GetLastError();
}


CSSAservice::~CSSAservice(void)
{
    if (m_hStoppedEvent)
    {
        CloseHandle(m_hStoppedEvent);
        m_hStoppedEvent = NULL;
    }
}


//
//   FUNCTION: CSSAservice::OnStart(DWORD, LPWSTR *)
//
//   PURPOSE: The function is executed when a Start command is sent to the 
//   service by the SCM or when the operating system starts (for a service 
//   that starts automatically). It specifies actions to take when the 
//   service starts. In this code sample, OnStart logs a service-start 
//   message to the Application log, and queues the main service function for 
//   execution in a thread pool worker thread.
//
//   PARAMETERS:
//   * dwArgc   - number of command line arguments
//   * lpszArgv - array of command line arguments
//
//   NOTE: A service application is designed to be long running. Therefore, 
//   it usually polls or monitors something in the system. The monitoring is 
//   set up in the OnStart method. However, OnStart does not actually do the 
//   monitoring. The OnStart method must return to the operating system after 
//   the service's operation has begun. It must not loop forever or block. To 
//   set up a simple monitoring mechanism, one general solution is to create 
//   a timer in OnStart. The timer would then raise events in your code 
//   periodically, at which time your service could do its monitoring. The 
//   other solution is to spawn a new thread to perform the main service 
//   functions, which is demonstrated in this code sample.
//
void CSSAservice::OnStart(DWORD dwArgc, LPWSTR *lpszArgv)
{
    // Log a service start message to the Application log.
    WriteEventLogEntry(TEXT("SSAservice Starting"));

    // Queue the main service function for execution in a worker thread.
    CThreadPool::QueueUserWorkItem(&CSSAservice::ServiceLoop, this);
}


//
//   FUNCTION: CSSAservice::ServiceLoop(void)
//
//   PURPOSE: The method performs the main function of the service. It runs 
//   on a thread pool worker thread.
//
void CSSAservice::ServiceLoop(void)
{
    const int nSeconds = 5;
    CheckDirectory();
    
    WIN32_FIND_DATA fd = { 0 };
    HANDLE hf = ::FindFirstFile(TEXT("*.ssa"), &fd);
    if (hf)
    {
      _agent = TEXT("SSAagent.exe");
      ::FindClose(hf);
    }
    else
      _agent = TEXT("Authoriz.exe");


    HANDLE hProcess = 0;
    DWORD pid = 0;
    FindOrCreateAgent(hProcess, pid);

    // Periodically check if the service is stopping.
    while (hProcess && !m_fStopping)
    {
      DWORD ret = ::WaitForSingleObject(hProcess, nSeconds * 1000);
      if (ret != WAIT_TIMEOUT)
      {
        WriteEventLogEntry(TEXT("Unexpected agent termination"), EVENTLOG_WARNING_TYPE);
        ::CloseHandle(hProcess);
        FindOrCreateAgent(hProcess, pid);
      }
    }

    if (hProcess)
    {
      ::CloseHandle(hProcess);
      HANDLE hAgent = ::OpenProcess(PROCESS_TERMINATE, FALSE, pid); 
      if (hAgent)
      {
        if (::TerminateProcess(hAgent, 0))
          _agent = NULL;
        else
          WriteErrorLogEntry(TEXT("TerminateProcess"));
        ::CloseHandle(hAgent);
      }
      else
        WriteErrorLogEntry(TEXT("OpenProcess(PROCESS_TERMINATE)"));
    }

    // Signal the stopped event.
    SetEvent(m_hStoppedEvent);
}


//
//   FUNCTION: CSSAservice::OnStop(void)
//
//   PURPOSE: The function is executed when a Stop command is sent to the 
//   service by SCM. It specifies actions to take when a service stops 
//   running. In this code sample, OnStop logs a service-stop message to the 
//   Application log, and waits for the finish of the main service function.
//
//   COMMENTS:
//   Be sure to periodically call ReportServiceStatus() with 
//   SERVICE_STOP_PENDING if the procedure is going to take long time. 
//
void CSSAservice::OnStop()
{
    // Log a service stop message to the Application log.
    WriteEventLogEntry(L"SSAservice Stopping");

    // Indicate that the service is stopping and wait for the finish of the 
    // main service function (ServiceWorkerThread).
    m_fStopping = TRUE;
    if (WaitForSingleObject(m_hStoppedEvent, INFINITE) != WAIT_OBJECT_0)
        throw GetLastError();
}