#define STRICT
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>
#include "service.hpp"
#include "message.h"

static CService * _pService;

CService::CService ()
{
  _pService = this;
}

BOOL  CService::ServicePause ()                 { return TRUE; };
BOOL  CService::ServiceContinue ()              { return TRUE; };
BOOL  CService::ServiceKill ()                  { return TRUE; };
DWORD CService::ServiceUserControl (DWORD)      { return SERVICE_RUNNING; };

void CService::Init (LPTSTR name, LPTSTR display, DWORD dwControls)
{
  bDebug = FALSE;
  lpServiceName = name;
  lpServiceDisplayName = display;
  dwControlsAccepted = (DWORD) dwControls;
  _pService = this;
}

void CService::SetStatusPending (DWORD dwWaitHint)
{
  if (bDebug)
    return;

  if (ssStatus.dwCurrentState == SERVICE_STOPPED || ssStatus.dwCurrentState == SERVICE_RUNNING)
    return;
        
  ssStatus.dwCheckPoint++;
  ssStatus.dwWaitHint = dwWaitHint;
  
  if (! SetServiceStatus (sshStatusHandle, &ssStatus))
    LogError(EVENTLOGMSG_SETSERVICESTATUS_ERROR);
}

void CService::ServiceControl (DWORD dwCtrlCode)
{
  DWORD dwRetVal;
  
  switch (dwCtrlCode)
    {
        
    case SERVICE_CONTROL_STOP:
      // give us some time
      ssStatus.dwCurrentState = SERVICE_STOP_PENDING;
      ssStatus.dwWaitHint = 1000;
      ssStatus.dwCheckPoint = 1;
      SetServiceStatus (sshStatusHandle, &ssStatus);

      dwRetVal = ServiceStop ();
      
      // report stop and exit code
      ssStatus.dwCurrentState = SERVICE_STOPPED;
      ssStatus.dwWaitHint = 0;
      ssStatus.dwCheckPoint = 0;
      ssStatus.dwServiceSpecificExitCode = 0;
      ssStatus.dwWin32ExitCode = dwRetVal;
      break;
      
    case SERVICE_CONTROL_PAUSE:
      // give us some time
      ssStatus.dwCurrentState = SERVICE_PAUSE_PENDING;
      ssStatus.dwWaitHint = 1000;
      ssStatus.dwCheckPoint = 1;
      SetServiceStatus (sshStatusHandle, &ssStatus);

      ServicePause ();
      
      // report pause success
      ssStatus.dwCurrentState = SERVICE_PAUSED;
      ssStatus.dwWaitHint = 0;
      ssStatus.dwCheckPoint = 0;
      break;
          
    case SERVICE_CONTROL_CONTINUE: 
      // give us some time
      ssStatus.dwCurrentState = SERVICE_STOP_PENDING;
      ssStatus.dwWaitHint = 1000;
      ssStatus.dwCheckPoint = 1;
      SetServiceStatus (sshStatusHandle, &ssStatus);

      ServiceContinue ();
      
      // report continue success
      ssStatus.dwCurrentState = SERVICE_RUNNING;
      ssStatus.dwWaitHint = 0;
      ssStatus.dwCheckPoint = 0;
      break;
            
    case SERVICE_CONTROL_SHUTDOWN:
      ssStatus.dwCurrentState = SERVICE_STOP_PENDING;
      ssStatus.dwWaitHint = 100; // shutdown should be quick
      ssStatus.dwCheckPoint = 1;
      SetServiceStatus (sshStatusHandle, &ssStatus);

      dwRetVal = ServiceKill ();
      
      ssStatus.dwCurrentState = SERVICE_STOPPED;
      ssStatus.dwWaitHint = 0;
      ssStatus.dwCheckPoint = 0;
      ssStatus.dwWin32ExitCode = dwRetVal;
      break;
      
    case SERVICE_CONTROL_INTERROGATE:  // update Status
      ssStatus.dwWaitHint = 0;
      ssStatus.dwCheckPoint = 0;
      break;

    default:
      if (dwCtrlCode >= SERVICE_USER_CONTROL && dwCtrlCode <= SERVICE_USER_CONTROL_LAST)
      {
        ssStatus.dwWaitHint = 0;
        ssStatus.dwCheckPoint = 0;
        ssStatus.dwServiceSpecificExitCode = 0;
        ssStatus.dwWin32ExitCode = 0;
        ssStatus.dwCurrentState = ServiceUserControl (dwCtrlCode);
      }
      break;
    }
  if (!SetServiceStatus (sshStatusHandle, &ssStatus))
    LogError(EVENTLOGMSG_SETSERVICESTATUS_ERROR);
}

VOID WINAPI CService::service_ctrl (DWORD dwCtrlCode)
{
  _pService->ServiceControl (dwCtrlCode);
}

VOID CService::service_main (DWORD dwArgc, LPTSTR * lpszArgv)
{
  _pService->ServiceMain (dwArgc, lpszArgv);
}

void CService::ServiceMain (DWORD dwArgc, LPTSTR * lpszArgv)
{
  sshStatusHandle = RegisterServiceCtrlHandler (TEXT (lpServiceName), service_ctrl);
  
  ssStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
  ssStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP|SERVICE_ACCEPT_SHUTDOWN;
  ssStatus.dwWin32ExitCode = 0;
  ssStatus.dwServiceSpecificExitCode = 0;
  ssStatus.dwCheckPoint = 0;
  ssStatus.dwWaitHint = 1000;
  
  if (sshStatusHandle)
  {
    ssStatus.dwCurrentState = SERVICE_START_PENDING;
    SetServiceStatus (sshStatusHandle, &ssStatus);

    if (ServiceStart (dwArgc, lpszArgv) == FALSE)
      ssStatus.dwCurrentState = SERVICE_STOPPED;
    else
      ssStatus.dwCurrentState = SERVICE_RUNNING;

    ssStatus.dwWaitHint = 0;
    if (!SetServiceStatus (sshStatusHandle, &ssStatus))
      LogError(EVENTLOGMSG_SETSERVICESTATUS_ERROR);
  }
  else
  {
    ssStatus.dwCurrentState = SERVICE_STOPPED;
    SetServiceStatus (sshStatusHandle, &ssStatus);
  }
  return;
}

void CService::Run(int argc, char **argv)
{
  char title[80];
  bDebug = (GetConsoleTitle (title, sizeof(title)) != 0);

  if (!bDebug)
  {
    SERVICE_TABLE_ENTRY dispatchTable[] =
    {
      {
        lpServiceName,  // name of the service
        service_main    // service function
      },
      {NULL, NULL}      // end of table
    };
  
    // the service control manager starts the service
    if (!StartServiceCtrlDispatcher (dispatchTable))
      LogError(EVENTLOGMSG_SERVICE_NOT_STARTED);
  }
  else
  {
    printf ("Console mode of '%s'.\n\n", lpServiceDisplayName);

    if (ServiceStart (argc, argv) == FALSE)
    {
      puts ("service start function fails");
      return;
    }

    printf("Service started\ncommands:\nstop, pause, continue, exit\n\n");
    
    for (;;)
    {
      char buf[80];
      printf("cmd:");
      gets(buf);

      if (strcmp(buf,"stop") == 0)
      {
        ServiceStop ();
        puts ("Service stopped");
        break;
      }
      else if (strcmp(buf,"pause") == 0)
      {
        ServicePause ();
        puts ("Service paused");
      }
      else if (strcmp(buf,"continue") == 0)
      {
        ServiceContinue ();
        puts ("Service continued");
      }
      else if (strcmp(buf,"exit") == 0)
      {
        ServiceKill ();
        puts ("Service killed");
        break;
      }
    }
  }
}

#define MAX_EVENT_STRINGS 8

DWORD CService::ReportServiceEvent(
    WORD EventType,
    DWORD EventId,
    DWORD SizeOfRawData,
    PVOID RawData,
    DWORD NumberOfStrings,
    ...
    )
{
    va_list arglist;
    HANDLE hEventSource;
    ULONG i;
    PWSTR Strings[ MAX_EVENT_STRINGS ];
    DWORD rv;

    hEventSource = RegisterEventSource (NULL, lpServiceName);
    if (hEventSource == NULL)
      return GetLastError();

    va_start( arglist, NumberOfStrings );

    if (NumberOfStrings > MAX_EVENT_STRINGS)
        NumberOfStrings = MAX_EVENT_STRINGS;

    for (i=0; i<NumberOfStrings; i++)
        Strings[ i ] = va_arg( arglist, PWSTR );


    if (!ReportEvent(   hEventSource,
                        EventType,
                        0,
                        EventId,
                        NULL,
                        (WORD)NumberOfStrings,
                        SizeOfRawData,
                        (LPCTSTR *) Strings,
                        RawData) )
        rv = GetLastError();
    else
        rv = ERROR_SUCCESS;

    DeregisterEventSource (hEventSource);
    return rv;
}

void CService::CmdInstallService ()
{
  char szPath[512];
  SC_HANDLE schService;
  SC_HANDLE schSCManager;

  GetModuleFileName (NULL, szPath, sizeof(szPath));

  schSCManager = OpenSCManager (NULL, NULL, SC_MANAGER_ALL_ACCESS);

  if (schSCManager)
  {
    schService = CreateService (
                      schSCManager,               // SCManager database
                      lpServiceName,              // name of service
                      lpServiceDisplayName,       // name to display
                      SERVICE_ALL_ACCESS,         // desired access
                      SERVICE_WIN32_OWN_PROCESS,  // service type
                      SERVICE_AUTO_START,         // start type
                      SERVICE_ERROR_NORMAL,       // error control type
                      szPath,                     // service's binary
                      NULL,                       // no load ordering group
                      NULL,                       // no tag identifier
                      "",                         // dependencies
                      NULL,                       // LocalSystem account
                      NULL);                      // no password

    if (schService)
    {
      printf ("%s installed.\n", lpServiceDisplayName);
      CloseServiceHandle (schService);
    }
    else
      printf ("CreateService failed\n");

    CloseServiceHandle (schSCManager);
  }
  else
    printf ("OpenSCManager failed - %s\n");


  // add registry values for event log
  HKEY hk; 
  DWORD dwData; 

  strcpy (szPath, "SYSTEM\\CurrentControlSet\\Services\\EventLog\\Application\\");
  strcat (szPath, lpServiceName);
   
  RegCreateKey(HKEY_LOCAL_MACHINE, szPath, &hk);

  GetModuleFileName(NULL, szPath, sizeof(szPath));

  RegSetValueEx(hk,               /* subkey handle         */ 
        "EventMessageFile",       /* value name            */ 
        0,                        /* must be zero          */ 
        REG_EXPAND_SZ,            /* value type            */ 
        (LPBYTE)szPath,           /* address of value data */ 
        strlen(szPath) + 1);      /* length of value data  */ 
 
  dwData = EVENTLOG_ERROR_TYPE | EVENTLOG_WARNING_TYPE | EVENTLOG_INFORMATION_TYPE;
 
  RegSetValueEx(hk,        /* subkey handle                */ 
        "TypesSupported",  /* value name                   */ 
        0,                 /* must be zero                 */ 
        REG_DWORD,         /* value type                   */ 
        (LPBYTE) &dwData,  /* address of value data        */ 
        sizeof(DWORD));    /* length of value data         */ 

  RegCloseKey(hk); 
}

void CService::CmdRemoveService ()
{
  SC_HANDLE schService;
  SC_HANDLE schSCManager;

  schSCManager = OpenSCManager (NULL, NULL, SC_MANAGER_ALL_ACCESS);

  if (schSCManager)
  {
    schService = OpenService (schSCManager, lpServiceName, SERVICE_ALL_ACCESS);

    if (schService)
    {
      // try to stop the service
      if (ControlService (schService, SERVICE_CONTROL_STOP, &ssStatus))
      {
        printf ("Stopping %s.", lpServiceDisplayName);
        Sleep (500);

        while (QueryServiceStatus (schService, &ssStatus))
        {
          if (ssStatus.dwCurrentState == SERVICE_STOP_PENDING)
          {
            printf (".");
            Sleep (500);
          }
          else
            break;
        }

        if (ssStatus.dwCurrentState == SERVICE_STOPPED)
          printf ("\n%s stopped.\n", lpServiceDisplayName);
        else
          printf ("\n%s failed to stop.\n", lpServiceDisplayName);

      }

      // now remove the service
      if (DeleteService (schService))
        printf ("%s removed.\n", lpServiceDisplayName);
      else
        printf ("DeleteService failed\n");


      CloseServiceHandle (schService);
    }
    else
      printf ("OpenService failed\n");

    CloseServiceHandle (schSCManager);
  }
  else
    printf ("OpenSCManager failed\n");
}

