zen `service` command

1 files changed, 308 insertions, 0 deletions

diff --git a/src/zen/cmds/service_cmd.cpp b/src/zen/cmds/service_cmd.cpp
new file mode 100644
index 000000000..a93bac37c
--- /dev/null
+++ b/src/zen/cmds/service_cmd.cpp
@@ -0,0 +1,308 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "service_cmd.h"
+
+#include <zencore/filesystem.h>
+#include <zencore/fmtutils.h>
+#include <zencore/logging.h>
+#include <zencore/zencore.h>
+#include <zenutil/service.h>
+
+#if ZEN_PLATFORM_WINDOWS
+#	include <zencore/windows.h>
+#	include <shellapi.h>
+#endif
+
+ZEN_THIRD_PARTY_INCLUDES_START
+#include <gsl/gsl-lite.hpp>
+ZEN_THIRD_PARTY_INCLUDES_END
+
+#include <memory>
+
+namespace zen {
+
+//////////////////////////////////////////////////////////////////////////
+#if ZEN_PLATFORM_WINDOWS
+BOOL
+WinIsElevated()
+{
+	BOOL   fRet	  = FALSE;
+	HANDLE hToken = NULL;
+	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
+	{
+		TOKEN_ELEVATION Elevation;
+		DWORD			cbSize = sizeof(TOKEN_ELEVATION);
+		if (GetTokenInformation(hToken, TokenElevation, &Elevation, sizeof(Elevation), &cbSize))
+		{
+			fRet = Elevation.TokenIsElevated;
+		}
+	}
+	if (hToken)
+	{
+		CloseHandle(hToken);
+	}
+	return fRet;
+}
+
+int
+WinRelaunchElevated()
+{
+	TCHAR CurrentDir[4096];
+	GetCurrentDirectory(4096, CurrentDir);
+
+	ExtendableWideStringBuilder<256> Parameters;
+	std::filesystem::path			 ExecutablePath = GetRunningExecutablePath();
+	std::wstring					 CommandLine(GetCommandLine());
+	std::string::size_type			 ExtraLength	  = CommandLine[0] == '\"' ? 2 : 0;
+	std::wstring					 CommandArguments = CommandLine.substr(ExecutablePath.string().length() + ExtraLength + 1);
+
+	ZEN_CONSOLE("Attempting to run '{} {}' elevated...", ExecutablePath, WideToUtf8(CommandArguments));
+
+	SHELLEXECUTEINFO shExInfo = {0};
+	shExInfo.cbSize			  = sizeof(shExInfo);
+	shExInfo.fMask			  = SEE_MASK_NOCLOSEPROCESS | SEE_MASK_NO_CONSOLE;
+	shExInfo.hwnd			  = 0;
+	shExInfo.lpVerb			  = TEXT("runas");			   // Operation to perform
+	shExInfo.lpFile			  = ExecutablePath.c_str();	   // Application to start
+	shExInfo.lpParameters	  = CommandArguments.c_str();  // Additional parameters
+	shExInfo.lpDirectory	  = CurrentDir;
+	shExInfo.nShow			  = SW_HIDE;
+	shExInfo.hInstApp		  = 0;
+
+	DWORD ReturnCode = 1;
+	if (ShellExecuteEx(&shExInfo))
+	{
+		WaitForSingleObject(shExInfo.hProcess, INFINITE);
+		GetExitCodeProcess(shExInfo.hProcess, &ReturnCode);
+		CloseHandle(shExInfo.hProcess);
+		if (ReturnCode == 0)
+		{
+			ZEN_CONSOLE("Elevated execution completed successfully.");
+		}
+		else
+		{
+			ZEN_CONSOLE("Elevated execution completed unsuccessfully, return code: '{}'.", ReturnCode);
+		}
+	}
+	else
+	{
+		ZEN_CONSOLE("Failed to elevated, operation did not complete.");
+	}
+	return (int)ReturnCode;
+}
+
+#endif
+
+ServiceCommand::ServiceCommand()
+{
+	m_Options.add_option("",
+						 "v",
+						 "verb",
+						 fmt::format("Verb for service - {}, {}, {}, {}, {}. Use '--' ",
+									 m_StatusOptions.program(),
+									 m_InstallOptions.program(),
+									 m_UninstallOptions.program(),
+									 m_StartOptions.program(),
+									 m_StopOptions.program()),
+						 cxxopts::value(m_Verb),
+						 "<verb>");
+	m_Options.parse_positional({"verb"});
+	m_Options.positional_help("verb");
+
+	m_StatusOptions.add_options()("h,help", "Print help");
+	m_StatusOptions.add_option("",
+							   "n",
+							   "name",
+							   fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							   cxxopts::value(m_ServiceName),
+							   "<name>");
+	m_StatusOptions.parse_positional({"name"});
+	m_StatusOptions.positional_help("name");
+
+	m_InstallOptions.add_options()("h,help", "Print help");
+	m_InstallOptions.add_option("", "s", "executable", "Path to server executable", cxxopts::value(m_ServerExecutable), "<path>");
+	m_InstallOptions.add_option("",
+								"n",
+								"name",
+								fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+								cxxopts::value(m_ServiceName),
+								"<name>");
+	m_InstallOptions.add_option("",
+								"d",
+								"display-name",
+								fmt::format("Service display, defaults to \"{}\"", m_ServiceDisplayName),
+								cxxopts::value(m_ServiceDisplayName),
+								"<display-name>");
+	m_InstallOptions.add_option("",
+								"",
+								"description",
+								fmt::format("Service description", m_ServiceDescription),
+								cxxopts::value(m_ServiceDescription),
+								"<description>");
+	m_InstallOptions.parse_positional({"executable", "name", "display-name"});
+	m_InstallOptions.positional_help("executable name display-name");
+
+	m_UninstallOptions.add_options()("h,help", "Print help");
+	m_UninstallOptions.add_option("",
+								  "n",
+								  "name",
+								  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+								  cxxopts::value(m_ServiceName),
+								  "<name>");
+	m_UninstallOptions.parse_positional({"name"});
+	m_UninstallOptions.positional_help("name");
+
+	m_StartOptions.add_options()("h,help", "Print help");
+	m_StartOptions.add_option("",
+							  "n",
+							  "name",
+							  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							  cxxopts::value(m_ServiceName),
+							  "<name>");
+	m_StartOptions.parse_positional({"name"});
+	m_StartOptions.positional_help("name");
+
+	m_StopOptions.add_options()("h,help", "Print help");
+	m_StopOptions.add_option("",
+							 "n",
+							 "name",
+							 fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							 cxxopts::value(m_ServiceName),
+							 "<name>");
+	m_StopOptions.parse_positional({"name"});
+	m_StopOptions.positional_help("name");
+}
+
+ServiceCommand::~ServiceCommand() = default;
+
+int
+ServiceCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
+{
+	ZEN_UNUSED(GlobalOptions);
+
+	using namespace std::literals;
+
+	std::vector<char*> SubCommandArguments;
+	cxxopts::Options*  SubOption			 = nullptr;
+	int				   ParentCommandArgCount = GetSubCommand(m_Options, argc, argv, m_SubCommands, SubOption, SubCommandArguments);
+	if (!ParseOptions(ParentCommandArgCount, argv))
+	{
+		return 0;
+	}
+
+	if (SubOption == nullptr)
+	{
+		throw zen::OptionParseException("command verb is missing");
+	}
+
+	if (!ParseOptions(*SubOption, gsl::narrow<int>(SubCommandArguments.size()), SubCommandArguments.data()))
+	{
+		return 0;
+	}
+
+	if (SubOption == &m_StatusOptions)
+	{
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Can't get information about service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE(
+			"Service '{}':\n"
+			"  Status:       {}\n"
+			"  Executable:   '{}'\n"
+			"  Display Name: '{}'\n"
+			"  Description:  '{}'",
+			m_ServiceName,
+			ToString(Info.Status),
+			Info.ExecutablePath,
+			Info.DisplayName,
+			Info.Description);
+	}
+
+	if (SubOption == &m_InstallOptions)
+	{
+#if ZEN_PLATFORM_WINDOWS
+		if (!WinIsElevated())
+		{
+			return WinRelaunchElevated();
+		}
+#endif	// ZEN_PLATFORM_WINDOWS
+		if (m_ServerExecutable.empty())
+		{
+			std::filesystem::path ExePath = zen::GetRunningExecutablePath();
+			ExePath.replace_filename("zenserver" + ExePath.extension().string());
+			m_ServerExecutable = ExePath;
+		}
+		m_ServerExecutable = std::filesystem::absolute(m_ServerExecutable);
+
+		std::error_code Ec = InstallService(m_ServerExecutable,
+											GlobalOptions.PassthroughCommandLine,
+											m_ServiceName,
+											m_ServiceDisplayName,
+											m_ServiceDescription);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to install service '{}' using '{}' . Reason: '{}'", m_ServiceName, m_ServerExecutable, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Installed service '{}' using '{}' successfully", m_ServiceName, m_ServerExecutable);
+	}
+
+	if (SubOption == &m_UninstallOptions)
+	{
+#if ZEN_PLATFORM_WINDOWS
+		if (!WinIsElevated())
+		{
+			return WinRelaunchElevated();
+		}
+#endif	// ZEN_PLATFORM_WINDOWS
+		std::error_code Ec = UninstallService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to uninstall service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Uninstalled service {} successfully", m_ServiceName);
+	}
+
+	if (SubOption == &m_StartOptions)
+	{
+#if ZEN_PLATFORM_WINDOWS
+		if (!WinIsElevated())
+		{
+			return WinRelaunchElevated();
+		}
+#endif	// ZEN_PLATFORM_WINDOWS
+		std::error_code Ec = StartService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Started service '{}' successfully", m_ServiceName);
+	}
+
+	if (SubOption == &m_StopOptions)
+	{
+#if ZEN_PLATFORM_WINDOWS
+		if (!WinIsElevated())
+		{
+			return WinRelaunchElevated();
+		}
+#endif	// ZEN_PLATFORM_WINDOWS
+		std::error_code Ec = StopService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to stop service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Stopped service '{}' successfully", m_ServiceName);
+	}
+
+	return 0;
+}
+
+}  // namespace zen