nssm.cpp

   1 #include "nssm.h"\r
   2 \r
   3 extern unsigned long tls_index;\r
   4 extern bool is_admin;\r
   5 extern imports_t imports;\r
   6 \r
   7 /* Are two strings case-insensitively equivalent? */\r
   8 int str_equiv(const TCHAR *a, const TCHAR *b) {\r
   9   size_t len = _tcslen(a);\r
  10   if (_tcslen(b) != len) return 0;\r
  11   if (_tcsnicmp(a, b, len)) return 0;\r
  12   return 1;\r
  13 }\r
  14 \r
  15 /* Convert a string to a number. */\r
  16 int str_number(const TCHAR *string, unsigned long *number, TCHAR **bogus) {\r
  17   if (! string) return 1;\r
  18 \r
  19   *number = _tcstoul(string, bogus, 0);\r
  20   if (**bogus) return 2;\r
  21 \r
  22   return 0;\r
  23 }\r
  24 \r
  25 int str_number(const TCHAR *string, unsigned long *number) {\r
  26   TCHAR *bogus;\r
  27   return str_number(string, number, &bogus);\r
  28 }\r
  29 \r
  30 /* Remove basename of a path. */\r
  31 void strip_basename(TCHAR *buffer) {\r
  32   size_t len = _tcslen(buffer);\r
  33   size_t i;\r
  34   for (i = len; i && buffer[i] != _T('\\') && buffer[i] != _T('/'); i--);\r
  35   /* X:\ is OK. */\r
  36   if (i && buffer[i - 1] == _T(':')) i++;\r
  37   buffer[i] = _T('\0');\r
  38 }\r
  39 \r
  40 /* How to use me correctly */\r
  41 int usage(int ret) {\r
  42   if (GetConsoleWindow()) print_message(stderr, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
  43   else popup_message(0, MB_OK, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
  44   return(ret);\r
  45 }\r
  46 \r
  47 void check_admin() {\r
  48   is_admin = false;\r
  49 \r
  50   /* Lifted from MSDN examples */\r
  51   PSID AdministratorsGroup;\r
  52   SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;\r
  53   if (! AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup)) return;\r
  54   CheckTokenMembership(0, AdministratorsGroup, /*XXX*/(PBOOL) &is_admin);\r
  55   FreeSid(AdministratorsGroup);\r
  56 }\r
  57 \r
  58 /* See if we were launched from a console window. */\r
  59 static void check_console() {\r
  60   /* If we're running in a service context there will be no console window. */\r
  61   HWND console = GetConsoleWindow();\r
  62   if (! console) return;\r
  63 \r
  64   unsigned long pid;\r
  65   if (! GetWindowThreadProcessId(console, &pid)) return;\r
  66 \r
  67   /*\r
  68     If the process associated with the console window handle is the same as\r
  69     this process, we were not launched from an existing console.  The user\r
  70     probably double-clicked our executable.\r
  71   */\r
  72   if (GetCurrentProcessId() != pid) return;\r
  73 \r
  74   /* We close our new console so that subsequent messages appear in a popup. */\r
  75   FreeConsole();\r
  76 }\r
  77 \r
  78 int num_cpus() {\r
  79   DWORD_PTR i, affinity, system_affinity;\r
  80   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
  81   for (i = 0; system_affinity & (1LL << i); i++);\r
  82   return (int) i;\r
  83 }\r
  84 \r
  85 int _tmain(int argc, TCHAR **argv) {\r
  86   check_console();\r
  87 \r
  88   /* Remember if we are admin */\r
  89   check_admin();\r
  90 \r
  91   /* Elevate */\r
  92   if (argc > 1) {\r
  93     /*\r
  94       Valid commands are:\r
  95       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
  96     */\r
  97     if (str_equiv(argv[1], _T("start"))) exit(control_service(0, argc - 2, argv + 2));\r
  98     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
  99     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 100     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 101     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 102     if (str_equiv(argv[1], _T("install"))) {\r
 103       if (! is_admin) {\r
 104         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL);\r
 105         exit(100);\r
 106       }\r
 107       exit(pre_install_service(argc - 2, argv + 2));\r
 108     }\r
 109     if (str_equiv(argv[1], _T("edit")) || str_equiv(argv[1], _T("get")) || str_equiv(argv[1], _T("set")) || str_equiv(argv[1], _T("reset")) || str_equiv(argv[1], _T("unset"))) {\r
 110       if (! is_admin) {\r
 111         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT);\r
 112         exit(100);\r
 113       }\r
 114       int ret = pre_edit_service(argc - 1, argv + 1);\r
 115       /* There might be a password here. */\r
 116       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 117       exit(ret);\r
 118     }\r
 119     if (str_equiv(argv[1], _T("remove"))) {\r
 120       if (! is_admin) {\r
 121         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE);\r
 122         exit(100);\r
 123       }\r
 124       exit(pre_remove_service(argc - 2, argv + 2));\r
 125     }\r
 126   }\r
 127 \r
 128   /* Thread local storage for error message buffer */\r
 129   tls_index = TlsAlloc();\r
 130 \r
 131   /* Register messages */\r
 132   if (is_admin) create_messages();\r
 133 \r
 134   /*\r
 135     Optimisation for Windows 2000:\r
 136     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 137     will time out after a few seconds on Windows 2000.  On newer versions the\r
 138     call returns instantly.  Check for stdin first and only try to call the\r
 139     function if there's no input stream found.  Although it's possible that\r
 140     we're running with input redirected it's much more likely that we're\r
 141     actually running as a service.\r
 142     This will save time when running with no arguments from a command prompt.\r
 143   */\r
 144   if (_fileno(stdin) < 0) {\r
 145     /* Set up function pointers. */\r
 146     if (get_imports()) exit(111);\r
 147 \r
 148     /* Start service magic */\r
 149     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 150     if (! StartServiceCtrlDispatcher(table)) {\r
 151       unsigned long error = GetLastError();\r
 152       /* User probably ran nssm with no argument */\r
 153       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 154       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 155       free_imports();\r
 156       exit(100);\r
 157     }\r
 158   }\r
 159   else exit(usage(1));\r
 160 \r
 161   /* And nothing more to do */\r
 162   exit(0);\r
 163 }\r