/*
Copyright (c) 2008-2010, 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 the Dirk Krause nor the names of
  other 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	irshdown.c	The irshdown program.
*/



#include "rshdown.h"

#include <windows.h>
#include <winbase.h>
#include <process.h>
#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <conio.h>

#include "rshdownm.h"




#line 59 "irshdown.ctr"




/**	Exit status.
*/
static int exval = 1;



/**	Command: Install.
*/
#define ACTION_INSTALL		0

/**	Command: Uninstall.
*/
#define ACTION_UNINSTALL	1


/**	Texts used by the program.
*/
static WCHAR *n[] = {
/*   0 */       L"\\rshdown.exe",       /* binary name */
/*   1 */       L"RshDown",             /* display name of service */
/*   2 */       L"Tcpip\0Dnscache\0",   /* service dependencies */
/*   3 */       L"",                    /* empty password for service account */
/*   4 */       L"Sender host (IP or host name): ", /* IP or host name */
/*   5 */       L"Sender port number: ",        /* sender port number */
/*   6 */       L"Local port number: ", /* local port number */
/*   7 */       L"Software\\RshDown",   /* registry key */
/*   8 */       L"sender",
/*   9 */       L"file",
/*  10 */       L"port",
/*  11 */       L"Shut down computer when receiving shutdown datagram",
/*  12 */       L"Datagram file: ",
/*  13 */       L"System\\CurrentControlSet\\Services\\RshDown",
/*  14 */       L"Description",
/*  15 */       L"\\rshdownm.dll",
/*  16 */       L"System\\CurrentControlSet\\Services\\Eventlog\\System\\RshDown",
/*  17 */       L"EventMessageFile",            /* string file */
/*  18 */       L"CategoryMessageFile",         /* string file */
/*  19 */       L"CategoryCount",               /* DWORD 1 */
/*  20 */       L"TypesSupported",              /* DWORD 7 */
NULL
};



/**	Options.
*/
WCHAR *options[] = {
L"-i",
L"--install",
L"-u",
L"--uninstall",
NULL
};



/**	Buffer for current working directory.
*/
static WCHAR b[1024];

/**	File name of executable file.
*/
static WCHAR f[SIZEOF(b,WCHAR)];

/**	Sender host name.
*/
static WCHAR h[1024];

/**	Datagram file name
*/
static WCHAR fn[1024];



/**	Find index of string \a s in array \a a.
	@param	a	Array of strings.
	@param	s	String to search.
	@return	Index of \a s in \a a or -1 (not found).
*/
int
array_index(WCHAR **a, WCHAR *s)
{
  int back = -1;
  WCHAR **ptr; int i;
  
  ptr = a; i = 0;
  while((*ptr != NULL) && (back == -1)) {
    if(wcscmp(*ptr, s) == 0) {
      back = i;
    }
    ptr++; i++;
  } 
  return back;
}



/**	Chomp a text (Remove leading and trailing whitespaces).
	@param	s	Original text.
	@return	Pointer to start of text or NULL.
*/
static
WCHAR *
do_chomp(WCHAR *s)
{
  WCHAR *back = NULL;
  WCHAR *trail = NULL;
  WCHAR *ptr;
  ptr = s;
  while(*ptr) {
    switch(*ptr) {
      case L' ': case L'\t': case L'\r': case L'\n': {
        if(trail == NULL) {
          trail = ptr;
        }
      } break;
      default: {
        if(back == NULL) {
          back = ptr;
        }
        trail = NULL;
      } break;
    }
    ptr++;
  }
  if(trail) {
    *trail = L'\0';
  }
  return back;
}



/**	Uninstall rshdown service.
*/
static
void
uninstall(void)
{
  BOOL removed;
  SC_HANDLE     hManager = NULL;
  SC_HANDLE     hService = NULL;
  SERVICE_STATUS        svcStatus;
  DWORD         dwCount;
  LONG          res;

  fprintf(stderr, "Going to uninstall the service.\n"); fflush(stderr);
  removed = FALSE;
  dwCount = 0;
  hManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
  if(hManager != NULL) {
    hService = OpenService(hManager, n[1], SERVICE_ALL_ACCESS);
    if(hService != NULL) {
      if(QueryServiceStatus(hService, &svcStatus)) {
        if(svcStatus.dwCurrentState != SERVICE_STOPPED) {
          if(ControlService(hService, SERVICE_CONTROL_STOP, &svcStatus)) {
            do {
              Sleep(1000);
              if(QueryServiceStatus(hService, &svcStatus)) {
                if(svcStatus.dwCurrentState == SERVICE_STOPPED) {
                  removed = TRUE;
                } else {
                  dwCount++;
                }
              } else {
                dwCount = 60;
              }
            } while((dwCount < 60) && (removed == FALSE));
            if(!removed) {
              fprintf(stderr, "ERROR: Failed to stop the service!\n");
              fflush(stderr);
            }
          }
        } else {
          removed = TRUE;
        }
        if(removed) {
          removed = FALSE;
          if(DeleteService(hService)) {
            removed = TRUE;
            res = RegDeleteKey(HKEY_LOCAL_MACHINE, n[16]);
            if(res != ERROR_SUCCESS) {
              fprintf(stderr, "ERROR: Failed to remove eventlog registry key!\n");
              fflush(stderr);
            }
            res = RegDeleteKey(HKEY_LOCAL_MACHINE, n[13]);
          } else {
            fprintf(stderr, "ERROR: Failed to uninstall the service!\n");
            fflush(stderr);
          }
        }
      } else {
        fprintf(stderr, "ERROR: Failed to query service state!\n");
        fflush(stderr);
      }
      CloseServiceHandle(hService);
    } else {
      fprintf(stderr, "ERROR: Service not found!\n");
      fflush(stderr);
    }
    CloseServiceHandle(hManager);
  } else {
    fprintf(stderr, "ERROR: Failed to connect to service manager!\n");
    fflush(stderr);
  }
}



/**	Install rshdown as a service.
	@return	TRUE on success, FALSE on error.
*/
static
BOOL
install_the_service(void)
{
  BOOL back = FALSE;
  SC_HANDLE hManager;
  SC_HANDLE hService;

  hManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
  if(hManager != NULL) {
    hService = CreateService(
      hManager,                         /* manager handle */
      n[1],                             /* service name */
      n[1],                             /* display name */
      SERVICE_ALL_ACCESS,               /* access rights to service */
      SERVICE_WIN32_OWN_PROCESS,        /* service type */
      SERVICE_AUTO_START,               /* SERVICE_DEMAND_START */
      SERVICE_ERROR_NORMAL,             /* error control */
      f,                                /* binary path name */
      NULL,                             /* load order group */
      NULL,                             /* tag ID */
      n[2],                             /* dependencies */
      NULL,                             /* service account name */
      n[3]                              /* password */
    );
    if(hService != NULL) {
      back = TRUE;
      CloseServiceHandle(hService);
    }
    CloseServiceHandle(hManager);
  } else {
    fprintf(stderr, "ERROR: Failed to connect to service manager!\n");
    fflush(stderr);
  }
  return back;
}



/**	Save expandable text value to registry.
	@param	hk	Registry key.
	@param	k	Entry name.
	@param	v	Entry value.
	@return	TRUE on success, FALSE on error.
*/
static
BOOL
saveKeyValueExpand(HKEY hk, WCHAR *k, WCHAR *v)
{
  BOOL back = FALSE;
  LONG res;
  res = RegSetValueEx(
    hk,
    k,
    0,
    REG_EXPAND_SZ,
    (const BYTE *)v,
    (sizeof(WCHAR)*(wcslen(v)+1))
  );
  if(res == ERROR_SUCCESS) {
    back = TRUE;
  }
  return back;
}



/**	Save text value to registry.
	@param	hk	Registry key.
	@param	k	Entry name.
	@param	v	Entry value.
	@return	TRUE on success, FALSE on error.
*/
static
BOOL
saveKeyValue(HKEY hk, WCHAR *k, WCHAR *v)
{
  BOOL back = FALSE;
  LONG res;
  res = RegSetValueEx(
    hk,
    k,
    0,
    REG_SZ,
    (const BYTE *)v,
    (sizeof(WCHAR)*(wcslen(v)+1))
  );
  if(res == ERROR_SUCCESS) {
    back = TRUE;
  }
  return back;
}



/**	Save DWORD value to registry.
	@param	hk	Registry key.
	@param	k	Entry name
	@param	v	DWORD value.
	@return	TRUE on success, FALSE on error.
*/
static
BOOL
saveDword(HKEY hk, WCHAR *k, DWORD v)
{
  BOOL back = FALSE;
  LONG res;
  res = RegSetValueEx(
    hk,
    k,
    0,
    REG_DWORD,
    (const BYTE *)(&v),
    sizeof(DWORD)
  );
  if(res == ERROR_SUCCESS) {
    back = TRUE;
  }
  return back;
}



/**	Install rshdown as a service.
*/
static
void
install(void)
{
  WCHAR il[1024], *ptr, *ptr2;
  BOOL back = FALSE;
  unsigned u;
  HKEY hk;
  DWORD dwDisp;
  LONG res;
  fprintf(stderr, "Going to install the service.\n"); fflush(stderr);
  if(GetModuleFileName(GetModuleHandle(NULL), b, SIZEOF(b,WCHAR))) {
    ptr = wcsrchr(b, L'\\');
    if(ptr) {
      *ptr = L'\0';
      wcscpy(f, b);
      if((wcslen(f) + wcslen(n[0])) < SIZEOF(f,WCHAR)) {
        wcscat(f, n[0]);
        fputws(n[4], stdout);           /* sender (IP/host) */
        fgetws(il, SIZEOF(il,WCHAR)-1, stdin);
        ptr = do_chomp(il);
        if(ptr) {
          wcscpy(h, il);
          fputws(n[5], stdout);         /* sender port */
          fgetws(il, SIZEOF(il,WCHAR)-1, stdin);
          if(swscanf(il, L"%u", &u) == 1) {
#if _MSC_VER >= 1400
	    swprintf(il, (sizeof(il)/sizeof(WCHAR)), L":%u", u);
#else
            swprintf(il, L":%u", u);
#endif
            if((wcslen(h)+wcslen(il)) < SIZEOF(h,WCHAR)) {
              wcscat(h, il);
              fputws(n[6], stdout);     /* local port */
              fgetws(il, SIZEOF(il,WCHAR)-1, stdin);
              if(swscanf(il, L"%u", &u) == 1) {
#if _MSC_VER >= 1400
		swprintf(il, (sizeof(il)/sizeof(WCHAR)), L"%u", u);
#else
                swprintf(il, L"%u", u);
#endif
                fputws(n[12], stdout);
                fgetws(fn, SIZEOF(fn,WCHAR)-1, stdin);
                ptr2 = do_chomp(fn);
                if(ptr2) {
                  res = RegCreateKeyEx(
                    HKEY_LOCAL_MACHINE,           /* parent key */
                    n[7],                         /* key name */
                    0,                            /* reserviced */
                    NULL,                         /* class */
                    REG_OPTION_NON_VOLATILE,      /* option */
                    KEY_ALL_ACCESS,               /* access permissions */
                    NULL,                         /* security */
                    &hk,                          /* result */
                    &dwDisp                       /* disposition */
                  );
                  if(res == ERROR_SUCCESS) {
                    if(saveKeyValue(hk, n[8], h)) {
                      if(saveKeyValue(hk, n[9], fn)) {
                        if(saveKeyValue(hk, n[10], il)) {
                          back = install_the_service();
                        } else {
                          fprintf(stderr, "ERROR: Failed to save local port number!\n");
                          fflush(stderr);
                        }
                      } else {
                        fprintf(stderr, "ERROR: Failed to save file name!\n");
                        fflush(stderr);
                      }
                    } else {
                      fprintf(stderr, "ERROR: Failed to save sender host!\n");
                      fflush(stderr);
                    }
                    RegCloseKey(hk);
                    if(back) {
                      res = RegCreateKeyEx(
                        HKEY_LOCAL_MACHINE, n[13], 0, NULL,
                        REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS,
                        NULL, &hk, &dwDisp
                      );
                      if(res == ERROR_SUCCESS) {
                        if(!saveKeyValue(hk, n[14], n[11])) {
                          back = FALSE;
                          fprintf(stderr, "ERROR: Failed to save description for service!\n");
                          fflush(stderr);
                        }
                        RegCloseKey(hk);
                      } else {
                        back = FALSE;
                        fprintf(stderr, "ERROR: Failed to open service registry key!\n");
                        fflush(stderr);
                      }
                    }
                  } else {
                    fprintf(stderr, "ERROR: Failed to open registry key for configuration!\n");
                    fflush(stderr);
                  }
                } else {
                  fprintf(stderr, "ERROR: File name must not be empty!\n");
                  fflush(stderr);
                }
              } else {
               fprintf(stderr, "ERROR: Local port must be numeric!\n");
               fflush(stderr);
              }
            } else {
              fprintf(stderr, "ERROR: Combination of host:port too long!\n");
              fflush(stderr);
            }
          } else {
            fprintf(stderr, "ERROR: Sender port must be numeric!\n");
            fflush(stderr);
          }
        }
      } else {
        fprintf(stderr, "ERROR: Executable name too long!\n");
        fflush(stderr);
      }
    } else {
      fprintf(stderr, "ERROR: No backslash in name of irshdown.exe!\n");
      fflush(stderr);
    }
  } else {
    fprintf(stderr, "ERROR: Failed to find installation directory!\n");
    fflush(stderr);
  }
  if(back) {
    wcscpy(f, b);
    if((wcslen(f) + wcslen(n[15])) < SIZEOF(f,WCHAR)) {
      wcscat(f, n[15]);
      res = RegCreateKeyEx(
        HKEY_LOCAL_MACHINE,
        n[16],
        0,
        NULL,
        REG_OPTION_NON_VOLATILE,
        KEY_ALL_ACCESS,
        NULL,
        &hk,
        &dwDisp
      );
      if(res == ERROR_SUCCESS) {
        if(saveKeyValueExpand(hk, n[17], f)) {
          if(saveKeyValueExpand(hk, n[18], f)) {
            if(saveDword(hk, n[19], 1)) {
              dwDisp = EVENTLOG_SUCCESS
                     | EVENTLOG_ERROR_TYPE
                     | EVENTLOG_WARNING_TYPE
                     | EVENTLOG_INFORMATION_TYPE;
              if(!saveDword(hk, n[20], dwDisp)) {
                back = FALSE;
                fprintf(stderr, "ERROR: Failed to save message types!\n");
                fflush(stderr);
              }
            } else {
              back = FALSE;
              fprintf(stderr, "ERROR: Failed to save category count!\n");
              fflush(stderr);
            }
          } else {
            back = FALSE;
            fprintf(stderr, "ERROR: Failed to save category message file name!\n");
            fflush(stderr);
          }
        } else {
          back = FALSE;
          fprintf(stderr, "ERROR: Failed to save event message file name!\n");
          fflush(stderr);
        }
        RegCloseKey(hk);
      } else {
        back = FALSE;
        fprintf(stderr, "ERROR: Failed to open registry for messages DLL!\n");
        fflush(stderr);
      }
    } else {
      back = FALSE;
      fprintf(stderr, "ERROR: Message DLL file name too long!\n");
      fflush(stderr);
    }
  }
  if(!back) {
    fprintf(stderr, "Errors occured, uninstalling the RshDown service.\n");
    fflush(stderr);
    uninstall();
  } else {
    exval = 0;
    fprintf(stderr, "The RshDown service was installed successfully.\n");
    fflush(stderr);
  }
}



/**	Run the required action (install or uninstall).
	@param	action	Flag for install or uninstall.
*/
static
void
run(int action)
{
  switch(action) {
    case ACTION_UNINSTALL: {
      uninstall();
    } break;
    default: {
      install();
    } break;
  }
}



/**	The main() function of the irshdown 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[])
{
  int action = ACTION_INSTALL;
  
#line 621 "irshdown.ctr"

  
  if(argc > 1) {
    switch(array_index(options, argv[1])) {
      case 2: case 3: {
        action = ACTION_UNINSTALL; 
      } break;
    }
  }
  run(action);
  
  
#line 632 "irshdown.ctr"

  exit(exval); return exval;
}