Merge pull request #29 from es3n1n/wsc-service-check

feat(loader): check for wsc service before starting

cxx-shared/cxx-shared.vcxitems

@@ -17,7 +17,7 @@
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\ctx.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\defer.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\ipc.hpp" />
-    <ClInclude Include="$(MSBuildThisFileDirectory)shared\names.hpp" />
+    <ClInclude Include="$(MSBuildThisFileDirectory)shared\strings.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\native.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\util.hpp" />
   </ItemGroup>

cxx-shared/shared/names.hpp

@@ -1,13 +0,0 @@
-#pragma once
-#include <string_view>
-
-namespace names {
-    constexpr std::string_view kProjectName = "defendnot";
-    constexpr std::string_view kRepoUrl = "https://github.com/es3n1n/defendnot";
-    constexpr std::string_view kDefaultAVName = "dnot.sh";
-
-    constexpr std::string_view kVictimProcess = "Taskmgr.exe";
-    constexpr std::string_view kDllName = "defendnot.dll";
-
-    constexpr std::string_view kVersion = "1.2.0";
-} // namespace names

cxx-shared/shared/native.hpp

@@ -23,7 +23,7 @@ namespace native {
 
         auto function = reinterpret_cast<Ty>(GetProcAddress(mod, function_name.data()));
         if (function == nullptr) {
-            throw std::runtime_error(std::format("unable to obtain {} from {}", module_name, function_name));
+            throw std::runtime_error(std::format("unable to obtain {} from {}", function_name, module_name));
         }
         return function;
     }

cxx-shared/shared/strings.hpp

@@ -0,0 +1,30 @@
+#pragma once
+#include "shared/util.hpp"
+#include <string_view>
+
+namespace strings {
+    constexpr std::string_view kProjectName = "defendnot";
+    constexpr std::string_view kRepoUrl = "https://github.com/es3n1n/defendnot";
+    constexpr std::string_view kVersion = "1.2.0";
+
+    constexpr std::string_view kDefaultAVName = "dnot.sh";
+
+    constexpr std::string_view kVictimProcess = "Taskmgr.exe";
+    constexpr std::string_view kDllName = "defendnot.dll";
+
+    constexpr std::string_view kWSCUnavailableError = /// !winserver
+        "Windows Security Center (WSC) is not available on this machine.\n"
+        "For more details, please refer to: https://github.com/es3n1n/defendnot/issues/25";
+
+    constexpr std::string_view kWSCUnavailableErrorWinServer = /// winserver
+        "Windows Security Center (WSC) is not available on this machine.\n"
+        "This typically occurs on Windows Server operating systems, which are not supported by this tool.\n"
+        "For more details, please refer to: https://github.com/es3n1n/defendnot/issues/17";
+
+    inline std::string_view wsc_unavailable_error() noexcept {
+        if (shared::is_winserver()) {
+            return kWSCUnavailableErrorWinServer;
+        }
+        return kWSCUnavailableError;
+    }
+} // namespace strings

cxx-shared/shared/util.hpp

@@ -27,4 +27,10 @@ namespace shared {
             freopen_s(reinterpret_cast<FILE**>(stderr), "CONOUT$", "w", stderr);
         });
     }
+
+    inline bool is_winserver() {
+        OSVERSIONINFOEXW osvi = {sizeof(osvi), 0, 0, 0, 0, {0}, 0, 0, 0, VER_NT_WORKSTATION};
+        const auto cond_mask = VerSetConditionMask(0, VER_PRODUCT_TYPE, VER_EQUAL);
+        return !VerifyVersionInfoW(&osvi, VER_PRODUCT_TYPE, cond_mask);
+    }
 } // namespace shared

defendnot-loader/core/autorun.cpp

@@ -1,7 +1,7 @@
 #include "core/core.hpp"
 
 #include "shared/ctx.hpp"
-#include "shared/names.hpp"
+#include "shared/strings.hpp"
 
 #include <memory>
 #include <print>
@@ -17,7 +17,7 @@
 
 namespace loader {
     namespace {
-        constexpr std::string_view kTaskName = names::kProjectName;
+        constexpr std::string_view kTaskName = strings::kProjectName;
 
         /// A very basic implementation, a lot of stuff is missing
         template <typename Ty>
@@ -164,7 +164,7 @@ namespace loader {
             principal->put_RunLevel(TASK_RUNLEVEL_HIGHEST);
 
             /// Info
-            reg_info->put_Author(bstr_t(names::kRepoUrl.data()));
+            reg_info->put_Author(bstr_t(strings::kRepoUrl.data()));
 
             /// Start even if we're on batteries
             settings->put_DisallowStartIfOnBatteries(VARIANT_FALSE);

defendnot-loader/core/core.hpp

@@ -24,4 +24,5 @@ namespace loader {
     [[nodiscard]] HANDLE inject(std::string_view dll_path, std::string_view proc_name);
     [[nodiscard]] bool add_to_autorun(AutorunType type);
     [[nodiscard]] bool remove_from_autorun();
+    void ensure_environment();
 } // namespace loader

defendnot-loader/core/ensure_environment.cpp

@@ -0,0 +1,26 @@
+#include "core/core.hpp"
+#include "shared/strings.hpp"
+#include "util/scm.hpp"
+
+#include <format>
+#include <stdexcept>
+
+namespace loader {
+    void ensure_environment() {
+        auto manager = scm::Manager();
+        if (!manager.valid()) [[unlikely]] {
+            throw std::runtime_error("Unable to open scm::Manager");
+        }
+
+        auto service = manager.get_service(L"wscsvc");
+        if (!service.valid() || !service.query_status()) {
+            throw std::runtime_error(strings::wsc_unavailable_error().data());
+        }
+
+        const auto state = service.state();
+        if (state != scm::ServiceState::RUNNING) {
+            throw std::runtime_error(std::format("{}\nService state: {}, but should be RUNNING", strings::wsc_unavailable_error(),
+                                                 scm::kServiceStateNames[static_cast<std::underlying_type_t<scm::ServiceState>>(state)]));
+        }
+    }
+} // namespace loader

defendnot-loader/defendnot-loader.vcxproj

@@ -250,11 +250,13 @@
   </ItemDefinitionGroup>
   <ItemGroup>
     <ClCompile Include="core\autorun.cpp" />
+    <ClCompile Include="core\ensure_environment.cpp" />
     <ClCompile Include="core\inject.cpp" />
     <ClCompile Include="main.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="core\core.hpp" />
+    <ClInclude Include="util\scm.hpp" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">

defendnot-loader/defendnot-loader.vcxproj.filters

@@ -24,10 +24,16 @@
     <ClCompile Include="core\autorun.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="core\ensure_environment.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="core\core.hpp">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="util\scm.hpp">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
 </Project>

defendnot-loader/main.cpp

@@ -2,7 +2,7 @@
 #include "shared/ctx.hpp"
 #include "shared/defer.hpp"
 #include "shared/ipc.hpp"
-#include "shared/names.hpp"
+#include "shared/strings.hpp"
 #include <argparse/argparse.hpp>
 
 #include <format>
@@ -38,12 +38,12 @@ namespace {
         std::println("** loading defendnot");
 
         auto dll_path = shared::get_this_module_path().parent_path();
-        dll_path /= names::kDllName;
+        dll_path /= strings::kDllName;
         if (!std::filesystem::exists(dll_path)) {
-            throw std::runtime_error(std::format("{} does not exist!", names::kDllName));
+            throw std::runtime_error(std::format("{} does not exist!", strings::kDllName));
         }
 
-        return loader::inject(dll_path.string(), names::kVictimProcess);
+        return loader::inject(dll_path.string(), strings::kVictimProcess);
     }
 
     void wait_for_finish(shared::InterProcessCommunication& ipc) {
@@ -65,8 +65,8 @@ namespace {
 
     void banner(const loader::Config& config) {
         std::println();
-        std::println("thanks for using {}", names::kProjectName);
-        std::println("please don't forget to leave a star at {}", names::kRepoUrl);
+        std::println("thanks for using {}", strings::kProjectName);
+        std::println("please don't forget to leave a star at {}", strings::kRepoUrl);
 
         if (!config.from_autorun && config.alloc_console) {
             system("pause");
@@ -75,7 +75,7 @@ namespace {
 } // namespace
 
 int main(int argc, char* argv[]) try {
-    argparse::ArgumentParser program(std::format("{}-loader", names::kProjectName), names::kVersion.data(), argparse::default_arguments::none);
+    argparse::ArgumentParser program(std::format("{}-loader", strings::kProjectName), strings::kVersion.data(), argparse::default_arguments::none);
 
     const auto fatal_print = [](const std::string_view str) -> void {
         shared::alloc_console();
@@ -94,11 +94,11 @@ int main(int argc, char* argv[]) try {
         .help("shows version and exits")
         .default_value(false)
         .implicit_value(true)
-        .action([&fatal_print](const auto& /*unused*/) -> void { fatal_print(std::format("{}-loader v{}", names::kProjectName, names::kVersion)); });
+        .action([&fatal_print](const auto& /*unused*/) -> void { fatal_print(std::format("{}-loader v{}", strings::kProjectName, strings::kVersion)); });
 
     /// defendnot-loader parameters:
-    program.add_argument("-n", "--name").help("av display name").default_value(std::string(names::kDefaultAVName)).nargs(1);
-    program.add_argument("-d", "--disable").help(std::format("disable {}", names::kProjectName)).default_value(false).implicit_value(true);
+    program.add_argument("-n", "--name").help("av display name").default_value(std::string(strings::kDefaultAVName)).nargs(1);
+    program.add_argument("-d", "--disable").help(std::format("disable {}", strings::kProjectName)).default_value(false).implicit_value(true);
     program.add_argument("-v", "--verbose").help("verbose logging").default_value(false).implicit_value(true);
     program.add_argument("--silent").help("do not allocate console").default_value(false).implicit_value(true);
     program.add_argument("--autorun-as-user").help("create autorun task as currently logged in user").default_value(false).implicit_value(true);
@@ -132,6 +132,8 @@ int main(int argc, char* argv[]) try {
     }
 
     setup_window(config);
+    loader::ensure_environment();
+
     setup_context(config);
 
     /// \todo @es3n1n: move this to a separate function and add move ctor for ipc
@@ -150,6 +152,7 @@ int main(int argc, char* argv[]) try {
 
     return EXIT_SUCCESS;
 } catch (std::exception& err) {
+    shared::alloc_console();
     std::println(stderr, "** fatal error: {}", err.what());
     system("pause");
     return EXIT_FAILURE;

defendnot-loader/util/scm.hpp

@@ -0,0 +1,113 @@
+#pragma once
+#include <array>
+#include <memory>
+#include <optional>
+#include <string_view>
+#include <Windows.h>
+
+namespace scm {
+    using SCHandleRaw = SC_HANDLE;
+    using SCHandle = std::unique_ptr<std::remove_pointer_t<SCHandleRaw>, decltype(&CloseServiceHandle)>;
+
+    inline SCHandle make_sc_handle(SCHandleRaw handle) {
+        return SCHandle(handle, CloseServiceHandle);
+    }
+
+    enum class ServiceState : std::uint8_t {
+        UNKNOWN = 0,
+        STOPPED,
+        STOP_PENDING,
+        START_PENDIND,
+        RUNNING,
+        PAUSED,
+        PAUSE_PENDING,
+        CONTINUE_PENDING,
+    };
+    /// \note @es3n1n: we should use magic_enum once we have more than one enum where we need to get value name
+    constexpr auto kServiceStateNames =
+        std::to_array<std::string_view>({"UNKNOWN", "STOPPED", "STOP_PENDING", "START_PENDING", "RUNNING", "PAUSED", "PAUSE_PENDING", "CONTINUE_PENDING"});
+
+    class Service {
+    public:
+        Service(SCHandleRaw handle): handle_(make_sc_handle(handle)) { };
+        ~Service() = default;
+
+    public:
+        bool query_status() noexcept {
+            if (status_process_.has_value()) {
+                return true;
+            }
+
+            SERVICE_STATUS_PROCESS status;
+            DWORD needed = 0;
+            if (!QueryServiceStatusEx(handle_.get(), SC_STATUS_PROCESS_INFO, reinterpret_cast<PBYTE>(&status), sizeof(status), &needed)) {
+                return false;
+            }
+
+            status_process_ = status;
+            return true;
+        }
+
+        [[nodiscard]] ServiceState state() noexcept {
+            if (!query_status() || !status_process_.has_value()) {
+                return ServiceState::UNKNOWN;
+            }
+
+            switch (status_process_->dwCurrentState) {
+            case SERVICE_STOPPED:
+                return ServiceState::STOPPED;
+            case SERVICE_STOP_PENDING:
+                return ServiceState::STOP_PENDING;
+            case SERVICE_START_PENDING:
+                return ServiceState::START_PENDIND;
+            case SERVICE_RUNNING:
+                return ServiceState::RUNNING;
+            case SERVICE_PAUSED:
+                return ServiceState::PAUSED;
+            case SERVICE_PAUSE_PENDING:
+                return ServiceState::PAUSE_PENDING;
+            case SERVICE_CONTINUE_PENDING:
+                return ServiceState::CONTINUE_PENDING;
+            default:
+                return ServiceState::UNKNOWN;
+            }
+        }
+
+        [[nodiscard]] bool valid() const noexcept {
+            return handle_.get() != nullptr;
+        }
+
+        [[nodiscard]] explicit operator bool() const noexcept {
+            return valid();
+        }
+
+    private:
+        std::optional<SERVICE_STATUS_PROCESS> status_process_ = std::nullopt;
+        SCHandle handle_;
+    };
+
+    class Manager {
+        constexpr static auto kDesiredPermissions = GENERIC_READ;
+
+    public:
+        Manager(): handle_(make_sc_handle(OpenSCManagerW(nullptr, nullptr, kDesiredPermissions))) { };
+        ~Manager() = default;
+
+    public:
+        [[nodiscard]] Service get_service(const std::wstring_view service_name) noexcept {
+            return Service(OpenServiceW(handle_.get(), service_name.data(), kDesiredPermissions));
+        }
+
+    public:
+        [[nodiscard]] bool valid() const noexcept {
+            return handle_.get() != nullptr;
+        }
+
+        [[nodiscard]] explicit operator bool() const noexcept {
+            return valid();
+        }
+
+    private:
+        SCHandle handle_;
+    };
+} // namespace scm

defendnot/core/com.hpp

@@ -5,6 +5,7 @@
 #include <thread>
 
 #include "core/log.hpp"
+#include "shared/strings.hpp"
 
 #include <Windows.h>
 
@@ -53,6 +54,7 @@ namespace defendnot {
 
         if (delayed) {
             /// Sleep for additional 15 seconds to let WSC proceed all previous requests
+            /// \todo @es3n1n: there should be a better way to handle this
             std::this_thread::sleep_for(std::chrono::seconds(15));
         }
 
@@ -94,9 +96,7 @@ namespace defendnot {
             IWscAVStatus* result = nullptr;
             const auto status = CoCreateInstance(detail::CLSID_IWscAVStatus, 0, 1, detail::IID_IWscAVStatus, reinterpret_cast<LPVOID*>(&result));
             if (status == REGDB_E_CLASSNOTREG) {
-                throw std::runtime_error("Windows Security Center (WSC) is not available on this machine.\n"
-                                         "This typically occurs on Windows Server operating systems, which are not supported by this tool.\n"
-                                         "For more details, please refer to: https://github.com/es3n1n/defendnot/issues/17");
+                throw std::runtime_error(strings::wsc_unavailable_error().data());
             }
 
             com_checked(status);