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 static inline void block(unsigned int a, short x, short y, unsigned long n) {\r
  86   HANDLE h = GetStdHandle(STD_OUTPUT_HANDLE);\r
  87   TCHAR s = _T(' ');\r
  88 \r
  89   unsigned long out;\r
  90   COORD c = { x, y };\r
  91   FillConsoleOutputAttribute(h, a, n, c, &out);\r
  92   FillConsoleOutputCharacter(h, s, n, c, &out);\r
  93 }\r
  94 \r
  95 static inline void R(short x, short y, unsigned long n) {\r
  96   block(BACKGROUND_RED | BACKGROUND_INTENSITY, x, y, n);\r
  97 }\r
  98 \r
  99 static inline void r(short x, short y, unsigned long n) {\r
 100   block(BACKGROUND_RED, x, y, n);\r
 101 }\r
 102 \r
 103 static inline void b(short x, short y, unsigned long n) {\r
 104   block(0, x, y, n);\r
 105 }\r
 106 \r
 107 void banner() {\r
 108   short y = 0;\r
 109 \r
 110   b(0, y, 80);\r
 111   y++;\r
 112 \r
 113   b(0, y, 80);\r
 114   y++;\r
 115 \r
 116   b(0, y, 80);\r
 117   y++;\r
 118 \r
 119   b(0, y, 80);\r
 120   y++;\r
 121 \r
 122   b(0, y, 80);\r
 123   r(18, y, 5); r(28, y, 4); r(41, y, 4); r(68, y, 1);\r
 124   R(6, y, 5); R(19, y, 4); R(29, y, 1); R(32, y, 3); R(42, y, 1); R(45, y, 3); R(52, y, 5); R(69, y, 4);\r
 125   y++;\r
 126 \r
 127   b(0, y, 80);\r
 128   r(8, y, 4); r(20, y, 1); r(28, y, 1); r(33, y, 3); r(41, y, 1); r(46, y, 3); r (57, y, 1);\r
 129   R(9, y, 2); R(21, y, 1); R(27, y, 1); R(34, y, 1); R(40, y, 1); R(47, y, 1); R(54, y, 3); R(68, y, 3);\r
 130   y++;\r
 131 \r
 132   b(0, y, 80);\r
 133   r(12, y, 1); r(20, y, 1); r(26, y, 1); r(34, y, 2); r(39, y, 1); r(47, y, 2); r(67, y, 2);\r
 134   R(9, y, 3); R(21, y, 1); R(27, y, 1); R(40, y, 1); R(54, y, 1); R(56, y, 2); R(67, y, 1); R(69, y, 2);\r
 135   y++;\r
 136 \r
 137   b(0, y, 80);\r
 138   r(9, y, 1); r(20, y, 1); r(26, y, 1); r (35, y, 1); r(39, y, 1); r(48, y, 1); r(58, y, 1);\r
 139   R(10, y, 3); R(21, y, 1); R(27, y, 1); R(40, y, 1); R(54, y, 1); R(56, y, 2); R(67, y, 1); R(69, y, 2);\r
 140   y++;\r
 141 \r
 142   b(0, y, 80);\r
 143   r(9, y, 1); r(56, y, 1); r(66, y, 2);\r
 144   R(11, y, 3); R(21, y, 1); R(26, y, 2); R(39, y, 2); R(54, y, 1); R(57, y, 2); R(69, y, 2);\r
 145   y++;\r
 146 \r
 147   b(0, y, 80);\r
 148   r(9, y, 1); r(26, y, 1); r(39, y, 1); r(59, y, 1);\r
 149   R(12, y, 3); R(21, y, 1); R(27, y, 2); R(40, y, 2); R(54, y, 1); R(57, y, 2); R(66, y, 1); R(69, y, 2);\r
 150   y++;\r
 151 \r
 152   b(0, y, 80);\r
 153   r(9, y, 1); r(12, y, 4); r(30, y, 1); r(43, y, 1); r(57, y, 1); r(65, y, 2);\r
 154   R(13, y, 2); R(21, y, 1); R(27, y, 3); R(40, y, 3); R(54, y, 1); R(58, y, 2); R(69, y, 2);\r
 155   y++;\r
 156 \r
 157   b(0, y, 80);\r
 158   r(9, y, 1); r(13, y, 4); r(27, y, 7); r(40, y, 7);\r
 159   R(14, y, 2); R(21, y, 1); R(28, y, 5); R(41, y, 5); R(54, y, 1); R(58, y, 2); R(65, y, 1); R(69, y, 2);\r
 160   y++;\r
 161 \r
 162   b(0, y, 80);\r
 163   r(9, y, 1); r(60, y, 1); r(65, y, 1);\r
 164   R(14, y, 3); R(21, y, 1); R(29, y, 6); R(42, y, 6); R(54, y, 1); R(58, y, 2); R(69, y, 2);\r
 165   y++;\r
 166 \r
 167   b(0, y, 80);\r
 168   r(9, y, 1); r(31, y, 1); r(44, y, 1); r(58, y, 1); r(64, y, 1);\r
 169   R(15, y, 3); R(21, y, 1); R(32, y, 4); R(45, y, 4); R(54, y, 1); R(59, y, 2); R(69, y, 2);\r
 170   y++;\r
 171 \r
 172   b(0, y, 80);\r
 173   r(9, y, 1); r(33, y, 1); r(46, y, 1); r(61, y, 1); r(64, y, 1);\r
 174   R(16, y, 3); R(21, y, 1); R(34, y, 2); R(47, y, 2); R(54, y, 1); R(59, y, 2); R(69, y, 2);\r
 175   y++;\r
 176 \r
 177   b(0, y, 80);\r
 178   r(9, y, 1); r(16, y, 4); r(36, y, 1); r(49, y, 1); r(59, y, 1); r(63, y, 1);\r
 179   R(17, y, 2); R(21, y, 1); R(34, y, 2); R(47, y, 2); R(54, y, 1); R(60, y, 2); R(69, y, 2);\r
 180   y++;\r
 181 \r
 182   b(0, y, 80);\r
 183   r(9, y, 1); r(17, y, 4); r(26, y, 1); r(36, y, 1); r(39, y, 1); r(49, y, 1);\r
 184   R(18, y, 2); R(21, y, 1); R(35, y, 1); R(48, y, 1); R(54, y, 1); R(60, y, 2); R(63, y, 1); R(69, y, 2);\r
 185   y++;\r
 186 \r
 187   b(0, y, 80);\r
 188   r(26, y, 2); r(39, y, 2); r(63, y, 1);\r
 189   R(9, y, 1); R(18, y, 4); R(35, y, 1); R(48, y, 1); R(54, y, 1); R(60, y, 3); R(69, y, 2);\r
 190   y++;\r
 191 \r
 192   b(0, y, 80);\r
 193   r(34, y, 1); r(47, y, 1); r(60, y, 1);\r
 194   R(9, y, 1); R(19, y, 3); R(26, y, 2); R(35, y, 1); R(39, y, 2); R(48, y, 1); R(54, y, 1); R(61, y, 2); R(69, y, 2);\r
 195   y++;\r
 196 \r
 197   b(0, y, 80);\r
 198   r(8, y, 1); r(35, y, 1); r(48, y, 1); r(62, y, 1); r(71, y, 1);\r
 199   R(9, y, 1); R(20, y, 2); R(26, y, 3); R(34, y, 1); R(39, y, 3); R(47, y, 1); R(54, y, 1); R(61, y, 1); R(69, y, 2);\r
 200   y++;\r
 201 \r
 202   b(0, y, 80);\r
 203   r(11, y, 1); r(26, y, 1); r(28, y, 5); r(39, y, 1); r(41, y, 5); r(51, y, 7); r(61, y, 1); r(66, y, 8);\r
 204   R(7, y, 4); R(21, y, 1); R(29, y, 1); R(33, y, 1); R(42, y, 1); R(46, y, 1); R(52, y, 5); R(67, y, 7);\r
 205   y++;\r
 206 \r
 207   b(0, y, 80);\r
 208   y++;\r
 209 \r
 210   b(0, y, 80);\r
 211   y++;\r
 212 }\r
 213 \r
 214 int _tmain(int argc, TCHAR **argv) {\r
 215   check_console();\r
 216 \r
 217 #ifdef UNICODE\r
 218   /*\r
 219     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
 220     mangled.  If we were compiled in ANSI mode it won't work.\r
 221    */\r
 222   _setmode(_fileno(stdout), _O_U16TEXT);\r
 223   _setmode(_fileno(stderr), _O_U16TEXT);\r
 224 #endif\r
 225 \r
 226   /* Remember if we are admin */\r
 227   check_admin();\r
 228 \r
 229   /* Elevate */\r
 230   if (argc > 1) {\r
 231     /*\r
 232       Valid commands are:\r
 233       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
 234     */\r
 235     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 236     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
 237     if (str_equiv(argv[1], _T("restart"))) {\r
 238       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
 239       if (ret) exit(ret);\r
 240       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 241     }\r
 242     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 243     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 244     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 245     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 246     if (str_equiv(argv[1], _T("install"))) {\r
 247       if (! is_admin) {\r
 248         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL);\r
 249         exit(100);\r
 250       }\r
 251       exit(pre_install_service(argc - 2, argv + 2));\r
 252     }\r
 253     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
 254       if (! is_admin) {\r
 255         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT);\r
 256         exit(100);\r
 257       }\r
 258       int ret = pre_edit_service(argc - 1, argv + 1);\r
 259       /* There might be a password here. */\r
 260       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 261       exit(ret);\r
 262     }\r
 263     if (str_equiv(argv[1], _T("remove"))) {\r
 264       if (! is_admin) {\r
 265         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE);\r
 266         exit(100);\r
 267       }\r
 268       exit(pre_remove_service(argc - 2, argv + 2));\r
 269     }\r
 270   }\r
 271 \r
 272   /* Thread local storage for error message buffer */\r
 273   tls_index = TlsAlloc();\r
 274 \r
 275   /* Register messages */\r
 276   if (is_admin) create_messages();\r
 277 \r
 278   /*\r
 279     Optimisation for Windows 2000:\r
 280     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 281     will time out after a few seconds on Windows 2000.  On newer versions the\r
 282     call returns instantly.  Check for stdin first and only try to call the\r
 283     function if there's no input stream found.  Although it's possible that\r
 284     we're running with input redirected it's much more likely that we're\r
 285     actually running as a service.\r
 286     This will save time when running with no arguments from a command prompt.\r
 287   */\r
 288   if (_fileno(stdin) < 0) {\r
 289     /* Set up function pointers. */\r
 290     if (get_imports()) exit(111);\r
 291 \r
 292     /* Start service magic */\r
 293     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 294     if (! StartServiceCtrlDispatcher(table)) {\r
 295       unsigned long error = GetLastError();\r
 296       /* User probably ran nssm with no argument */\r
 297       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 298       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 299       free_imports();\r
 300       exit(100);\r
 301     }\r
 302   }\r
 303   else exit(usage(1));\r
 304 \r
 305   /* And nothing more to do */\r
 306   exit(0);\r
 307 }\r