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 int num_cpus() {\r
  59   DWORD_PTR i, affinity, system_affinity;\r
  60   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
  61   for (i = 0; system_affinity & (1LL << i); i++);\r
  62   return (int) i;\r
  63 }\r
  64 \r
  65 int _tmain(int argc, TCHAR **argv) {\r
  66   check_console();\r
  67 \r
  68 #ifdef UNICODE\r
  69   /*\r
  70     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
  71     mangled.  If we were compiled in ANSI mode it won't work.\r
  72    */\r
  73   _setmode(_fileno(stdout), _O_U16TEXT);\r
  74   _setmode(_fileno(stderr), _O_U16TEXT);\r
  75 #endif\r
  76 \r
  77   /* Remember if we are admin */\r
  78   check_admin();\r
  79 \r
  80   /* Set up function pointers. */\r
  81   if (get_imports()) exit(111);\r
  82 \r
  83   /* Elevate */\r
  84   if (argc > 1) {\r
  85     /*\r
  86       Valid commands are:\r
  87       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
  88     */\r
  89     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
  90     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
  91     if (str_equiv(argv[1], _T("restart"))) {\r
  92       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
  93       if (ret) exit(ret);\r
  94       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
  95     }\r
  96     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
  97     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
  98     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
  99     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 100     if (str_equiv(argv[1], _T("install"))) {\r
 101       if (! is_admin) {\r
 102         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL);\r
 103         exit(100);\r
 104       }\r
 105       exit(pre_install_service(argc - 2, argv + 2));\r
 106     }\r
 107     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
 108       if (! is_admin) {\r
 109         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT);\r
 110         exit(100);\r
 111       }\r
 112       int ret = pre_edit_service(argc - 1, argv + 1);\r
 113       /* There might be a password here. */\r
 114       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 115       exit(ret);\r
 116     }\r
 117     if (str_equiv(argv[1], _T("remove"))) {\r
 118       if (! is_admin) {\r
 119         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE);\r
 120         exit(100);\r
 121       }\r
 122       exit(pre_remove_service(argc - 2, argv + 2));\r
 123     }\r
 124   }\r
 125 \r
 126   /* Thread local storage for error message buffer */\r
 127   tls_index = TlsAlloc();\r
 128 \r
 129   /* Register messages */\r
 130   if (is_admin) create_messages();\r
 131 \r
 132   /*\r
 133     Optimisation for Windows 2000:\r
 134     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 135     will time out after a few seconds on Windows 2000.  On newer versions the\r
 136     call returns instantly.  Check for stdin first and only try to call the\r
 137     function if there's no input stream found.  Although it's possible that\r
 138     we're running with input redirected it's much more likely that we're\r
 139     actually running as a service.\r
 140     This will save time when running with no arguments from a command prompt.\r
 141   */\r
 142   if (! GetStdHandle(STD_INPUT_HANDLE)) {\r
 143     /* Start service magic */\r
 144     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 145     if (! StartServiceCtrlDispatcher(table)) {\r
 146       unsigned long error = GetLastError();\r
 147       /* User probably ran nssm with no argument */\r
 148       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 149       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 150       free_imports();\r
 151       exit(100);\r
 152     }\r
 153   }\r
 154   else exit(usage(1));\r
 155 \r
 156   /* And nothing more to do */\r
 157   exit(0);\r
 158 }\r