/* ********************************************************************* */
/* *                                                                   * */
/* * Copyright (c) 2010 - Dipl.-Ing. Dirk Krause                       * */
/* *                                                                   * */
/* * All rights reserved.                                              * */
/* *                                                                   * */
/* * Redistribution and use in source and binary forms,                * */
/* * with or without modification, are permitted provided              * */
/* * that the following conditions are met:                            * */
/* *                                                                   * */
/* * * Redistributions of source code must retain the above            * */
/* *   copyright notice, this list of conditions and the               * */
/* *   following disclaimer.                                           * */
/* * * Redistributions in binary form must reproduce the above         * */
/* *   opyright notice, this list of conditions and the following      * */
/* *   disclaimer in the documentation and/or other materials          * */
/* *   provided with the distribution.                                 * */
/* * * Neither the name of Dirk Krause nor the names of                * */
/* *   contributors may be used to endorse or promote                  * */
/* *   products derived from this software without specific            * */
/* *   prior written permission.                                       * */
/* *                                                                   * */
/* * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND            * */
/* * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,       * */
/* * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF          * */
/* * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE          * */
/* * DISCLAIMED.                                                       * */
/* * IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE          * */
/* * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,             * */
/* * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT           * */
/* * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;          * */
/* * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)          * */
/* * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN         * */
/* * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE         * */
/* * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS           * */
/* * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH              * */
/* * DAMAGE.                                                           * */
/* *                                                                   * */
/* ********************************************************************* */

/**	@file	kwindown.c	The kwindown service.
*/



#include "kwintool.h"

#include "kwindownm.h"




#line 52 "kwindown.ctr"




/**	Maximum number of attempts to stop the service.
*/
static int max_attempts_to_stop = 5;

/**	Number of milliseconds to sleep between service stop attempts.
*/
static DWORD milli_seconds_to_stop = 200;

/**	Number of millicseconds to sleep between status requests.
*/
static DWORD milli_seconds_before_next_status = 300;

/**	Number of milliseconds to sleep between registry checks.
*/
static DWORD milli_seconds_in_service = 3000;

/**	Number of seconds to sleep if shutdown time not yet reached.
*/

static DWORD seconds_between_attempts = 60;

/**	Grace period during shutdown.
*/
static DWORD shutdown_grace_period = 30;


/**	Service name.
*/
static WCHAR service_name[]	= { L"kWinDown" };

/**	Parent registry key (non-volatile).
*/
static WCHAR parent_key[]	= { L"Software\\DKrause\\Shared" };

/**	Our registry key (volatile).
*/
static WCHAR temp_key[]		= { L"Software\\DKrause\\Shared\\kWinDown" };

/**	Name of registry entry.
*/
static WCHAR shutdown_time[]	= { L"ShutdownTimestamp" };

/**	Privilege name required for shutdown.
*/
static WCHAR privilege_name[]	= { L"SeShutdownPrivilege" };

/**	Message shown on screen for shutdown.
*/
static WCHAR shutdown_reason[]	= { L"System was unused for too long." };

/**	Message shown for "no shutdown".
*/
static WCHAR shutdown_time_no_shutdown[] = { L"No shutdown." };

/**	Service status.
*/
SERVICE_STATUS		service_status;

/**	Service handle.
*/
SERVICE_STATUS_HANDLE	h_service_status;

/**	Event handle: Must stop service (main thread to service thread).
*/
static volatile	HANDLE	hEventMustStop	= NULL;

/**	Event handle: Did stop service (Service thread to main thread).
*/
static volatile	HANDLE	hEventDidStop	= NULL;



/**	Log one message to event log system.
	@param	wt	Catalog.
	@param	eid	Event ID.
	@param	arg	Message arguments
*/
static void log_msg(WORD wt, DWORD eid, WCHAR *arg)
{
  HANDLE hEventLog;
  WCHAR *msgarray[2];
  hEventLog = RegisterEventSource(NULL, service_name);
  if(arg) {
    msgarray[0] = arg; msgarray[1] = NULL;
  } else {
    msgarray[0] = msgarray[1] = NULL;
  }
  if(hEventLog != NULL) {
    ReportEvent(
      hEventLog, 
      wt,
      CAT_KWINDOWN,
      eid,
      NULL,
      (arg ? 1 : 0),
      0,
      (arg ? msgarray : NULL),
      NULL
    );
    DeregisterEventSource(hEventLog); hEventLog = NULL;
  }
}



/**	Flag: The missing parent registry key was reported to event log.
*/
static int missing_parent_key_reported = 0;

/**	Flag: The missing registry key was reported to event log.
*/
static int missing_key_reported = 0;

/**	Flag: The wrong entry type was reported to event log.
*/
static int wrong_type_reported = 0;



/**	Get shutdown timestamp from registry.
	@param	sht	Pointer to shutdown timestamp variable (in+out).
*/
static
void
get_shutdown_time(__time64_t *sht)
{
  __time64_t old_shutdown_time;
  __time64_t back = (__time64_t)0UL;
  LONG res;
  DWORD dwDisp = 0, dwType = 0, dwSz = 0;
  HKEY hk;
  __time64_t myvalue;
  
  myvalue = (__time64_t)0UL;
  old_shutdown_time = *sht;
  res = RegCreateKeyEx(
    HKEY_LOCAL_MACHINE, parent_key, 0, NULL, REG_OPTION_NON_VOLATILE,
    KEY_ALL_ACCESS, NULL, &hk, &dwDisp
  );
  if(res == ERROR_SUCCESS) {    
    RegCloseKey(hk);
    res = RegCreateKeyEx(
      HKEY_LOCAL_MACHINE, temp_key, 0, NULL, REG_OPTION_VOLATILE,
      KEY_ALL_ACCESS, NULL, &hk, &dwDisp
    );
    if(res == ERROR_SUCCESS) {  
      dwType = REG_QWORD; dwSz = sizeof(__time64_t);
      res = RegQueryValueEx(
        hk, shutdown_time, NULL, &dwType, (void *)(&myvalue), &dwSz
      );
      if(res == ERROR_SUCCESS) {
        if(dwType == REG_QWORD) {
          if(dwSz == sizeof(__time64_t)) {
            back = myvalue;	
          } else {              
	    if(!wrong_type_reported) {
	      wrong_type_reported = 1;
	      log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_REGENTRY_WRONG_TYPE, NULL);
	    }
          }
        } else {                
	  if(!wrong_type_reported) {
	    wrong_type_reported = 1;
	    log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_REGENTRY_WRONG_TYPE, NULL);
	  }
        }
      } else {                  
        /* Value does not exist = 0 */
      }
      RegCloseKey(hk);
    } else {                    
      if(!missing_key_reported) {
        missing_key_reported = 1;
        log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_REGKEY_MISSING, temp_key);
      }
    }
  } else {                      
    if(!missing_parent_key_reported) {
      missing_parent_key_reported = 1;
      log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_REGKEY_MISSING, parent_key);
    }
  } 
  if(back != old_shutdown_time) {
    WCHAR buffer_old_time[64];
    struct tm *tm;
    if(back) {
      tm = _localtime64(&back);
      if(tm) {
        wsprintf(
          buffer_old_time,
	  L"%04d-%02d-%02d %02d:%02d:%02d",
	  (1900 + tm->tm_year),
	  (1 + tm->tm_mon),
	  tm->tm_mday,
	  tm->tm_hour,
	  tm->tm_min,
	  tm->tm_sec
        );
        log_msg(EVENTLOG_INFORMATION_TYPE, MSG_INFO_SHUTDOWN_TIME_CHANGED, buffer_old_time);
      }
    } else {
      log_msg(EVENTLOG_INFORMATION_TYPE, MSG_INFO_SHUTDOWN_TIME_CHANGED, shutdown_time_no_shutdown);
    }
  }
  *sht = back;
}



/**	Service thread function.
	@param	dataforthread	Additional data for thread (ignored).
*/
static void __cdecl
thread_function(void *dataforthread)
{
  /* LONG res;
 */
  /* HKEY hk;
 */
  /* DWORD dwDisp;
 */
  /* DWORD dwType;
 */
  /* DWORD dwSz;
 */
  /* DWORD dwValue;
 */
  DWORD waitResult;
  unsigned short must_shutdown = 0;
  int can_continue = 1;
  HANDLE	current_proc;
  HANDLE	proc_token;
  TOKEN_PRIVILEGES	ns;
  LUID		myluid;
  __time64_t	t_shutdown = (__time64_t)0UL, t_current = (__time64_t)0UL, t_last_attempt = (__time64_t)0UL;

  
#line 286 "kwindown.ctr"

  
  t_shutdown = (__time64_t)0UL;
  t_last_attempt = (__time64_t)0UL;
  can_continue = 1;
  log_msg(EVENTLOG_SUCCESS, MSG_KWINDOWN_UP, NULL);
  /*	Run until either shutdown timestamp reached or service stop request.
  */
  while(can_continue) {
    _time64(&t_current);
    if(t_last_attempt) {	/* Fheck registry in large intervals only. */
      if(t_current >= (t_last_attempt + (__time64_t)seconds_between_attempts))
      {
        t_last_attempt = t_current;
        get_shutdown_time(&t_shutdown);
      }
    } else {			/* First registry check after start. */
      t_last_attempt = t_current;
      get_shutdown_time(&t_shutdown);
    }
    if(t_shutdown) {
      if(t_current > t_shutdown) {
        t_last_attempt = t_current;
        get_shutdown_time(&t_shutdown);
        if(t_shutdown) {
          if(t_current > t_shutdown) {	/* Shutdown timestamp reached. */
            can_continue = 0; must_shutdown = 1;
          }
        }
      }
    }
    /* check event
 */
    if(can_continue) {
      waitResult =
      WaitForSingleObjectEx(hEventMustStop, milli_seconds_in_service, FALSE);
      switch(waitResult) {
        case WAIT_OBJECT_0: {   
          can_continue = 0;
        } break;
        default: {
        } break;
      }
    }
  }
  /*	Set event, shut down if necessary.
  */
  log_msg(EVENTLOG_INFORMATION_TYPE, MSG_INFO_SERVICE_DOWN, NULL);
  SetEvent(hEventDidStop);
  if(must_shutdown) {			
    log_msg(EVENTLOG_INFORMATION_TYPE, MSG_INFO_MUST_SHUTDOWN, NULL);
    /* ACTION: Start shutdown
 */
    current_proc = GetCurrentProcess();
    if(current_proc != (HANDLE)0) {	
      LookupPrivilegeValue(NULL, privilege_name, &myluid);
      if(OpenProcessToken(current_proc, (TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY),&proc_token)) {
        ns.PrivilegeCount = 1;		
        ns.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        ns.Privileges[0].Luid.HighPart = myluid.HighPart;
        ns.Privileges[0].Luid.LowPart = myluid.LowPart;
        if(AdjustTokenPrivileges(proc_token, FALSE, &ns, 0, NULL, NULL)) {
	  
        } else {			
          log_msg(
	    EVENTLOG_ERROR_TYPE,
	    MSG_ERROR_FAILED_TO_OBTAIN_SHUTDOWN_PRIVILEGE,
	    NULL
	  );
        }
        CloseHandle(proc_token);
      } else {				
        log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_ACCESS_TOKEN, NULL);
      }
    } else {				
      log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_PROCESS_HANDLE, NULL);
    }
    if(InitiateSystemShutdownEx(
      NULL, shutdown_reason, shutdown_grace_period, TRUE, FALSE,
      (SHTDN_REASON_MAJOR_OTHER | SHTDN_REASON_MINOR_OTHER
       | SHTDN_REASON_FLAG_PLANNED)
    ))
    {					
      log_msg(EVENTLOG_SUCCESS, MSG_SUCCESS_SHUTDOWN_INITIATED, NULL);
    } else {				
      log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_SHUTDOWN, NULL);
    }
  } else {
    log_msg(EVENTLOG_INFORMATION_TYPE, MSG_INFO_NO_SHUTDOWN, NULL);
  }
  
  
#line 375 "kwindown.ctr"

  _endthread();
}



/**	Service control event handler.
	@param	ctrl_code	Event type code.
*/
VOID
CALLBACK
service_ctrl_handler(IN DWORD ctrl_code)
{
  int i = 0, success = 0;
  DWORD waitResult;

  service_status.dwCheckPoint = 0;
  service_status.dwWaitHint = 0;
  switch(ctrl_code) {
    case SERVICE_CONTROL_SHUTDOWN: case SERVICE_CONTROL_STOP: {
      service_status.dwCurrentState = SERVICE_STOP_PENDING;
      service_status.dwWaitHint = milli_seconds_before_next_status;
      SetServiceStatus(h_service_status, &service_status);
      if((hEventMustStop != NULL) && (hEventDidStop != NULL)) {
        SetEvent(hEventMustStop);
        for(i = 0; i < max_attempts_to_stop; i++) {
          waitResult =
	  WaitForSingleObjectEx(hEventDidStop, milli_seconds_to_stop, FALSE);
	  switch(waitResult) {
	    case WAIT_ABANDONED: 		/* threads problem, should not happen */
	    case WAIT_IO_COMPLETION: {		/* I/O operation, should not happen */
	      i = max_attempts_to_stop;
	    } break;
	    case WAIT_OBJECT_0: {		/* signal was set */
	      i = max_attempts_to_stop;
	      ResetEvent(hEventDidStop);
	      success = 1;
	    } break;
	    case WAIT_TIMEOUT: {		/* timeout, continue */
	      service_status.dwCheckPoint++;
	      SetServiceStatus(h_service_status, &service_status);
	    } break;
	  }
        }
	if(!success) {
	  log_msg(
	    EVENTLOG_ERROR_TYPE,
	    MSG_ERROR_TIMEOUT_WHILE_WAITING_FOR_BGTHREAD, NULL
	  );
	}
      }
      service_status.dwCurrentState = SERVICE_STOPPED;
      service_status.dwCheckPoint = 0;
      service_status.dwWaitHint = 0;
      if(success) {
        if(hEventMustStop != NULL) {
          CloseHandle(hEventMustStop); hEventMustStop = NULL;
        }
        if(hEventDidStop != NULL) {
          CloseHandle(hEventDidStop); hEventDidStop = NULL;
        }
      }
    } break;
  }
  SetServiceStatus(h_service_status, &service_status);
}



/**	Service function.
	@param	argc	Number of command line arguments.
	@param	argv	Command line arguments array.
*/
VOID
CALLBACK
service_function(DWORD argc, LPTSTR *argv)
{
  uintptr_t thread_id;
  memset((void *)(&service_status), 0, sizeof(SERVICE_STATUS));
  service_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
  service_status.dwCurrentState = SERVICE_START_PENDING;
  service_status.dwControlsAccepted =
  SERVICE_ACCEPT_SHUTDOWN | SERVICE_ACCEPT_STOP;
  service_status.dwWin32ExitCode = 0;
  service_status.dwServiceSpecificExitCode = 0;
  service_status.dwCheckPoint = 0;
  service_status.dwWaitHint = 0;
  h_service_status =
  RegisterServiceCtrlHandler(service_name, service_ctrl_handler);
  if(h_service_status != (SERVICE_STATUS_HANDLE)0) {
    hEventMustStop = CreateEvent(NULL, TRUE, FALSE, NULL);
    if(hEventMustStop != NULL) {
      hEventDidStop = CreateEvent(NULL, TRUE, FALSE, NULL);
      if(hEventDidStop != NULL) {
        ResetEvent(hEventMustStop);
        ResetEvent(hEventDidStop);
	thread_id = _beginthread(thread_function, 0, NULL);
	if(thread_id != -1L) {
	  service_status.dwCurrentState = SERVICE_RUNNING;
	} else {
	  service_status.dwCurrentState = SERVICE_STOPPED;
	  log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_START_THREAD, NULL);
	  SetEvent(hEventDidStop);
	}
      } else {
        service_status.dwCurrentState = SERVICE_STOPPED;
        log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_CREATE_EVENT, NULL);
      }
    } else {
      service_status.dwCurrentState = SERVICE_STOPPED;
      log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_CREATE_EVENT, NULL);
    }
    SetServiceStatus(h_service_status, &service_status);
  } else {
    log_msg(EVENTLOG_ERROR_TYPE, MSG_ERROR_FAILED_TO_REGISTER_SERVICE, NULL);
  }
}



/**	Connect service names to service functions.
*/
static SERVICE_TABLE_ENTRY dispatch_table[]= {
  { service_name, service_function },
  { NULL, NULL }
};



/**	Main program.
	@param	argc	Number of command line arguments.
	@param	argv	Command line arguments array.
	@return	0 on success, any other value indicates an error.
*/
int wmain(int argc, WCHAR *argv[])
{
  if(StartServiceCtrlDispatcher(dispatch_table) == 0) {
    printf("ERROR: This program is a Windows service and can not be used\n");
    printf("on the console.\n");
  }
  exit(0); return 0;
}

