Merge pull request #32 from es3n1n/update-substatuses

feat(bootstrap): set scan/settings/protupdate substatuses to no_action

cxx-shared/cxx-shared.vcxitems

@@ -14,6 +14,7 @@
     <ProjectCapability Include="SourceItemsFromImports" />
   </ItemGroup>
   <ItemGroup>
+    <ClInclude Include="$(MSBuildThisFileDirectory)shared\com.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\ctx.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\defer.hpp" />
     <ClInclude Include="$(MSBuildThisFileDirectory)shared\ipc.hpp" />

cxx-shared/shared/com.hpp

@@ -0,0 +1,121 @@
+#pragma once
+#include <Windows.h>
+
+#include "shared/strings.hpp"
+
+#include <cstdint>
+#include <format>
+#include <print>
+#include <source_location>
+#include <stdexcept>
+
+namespace com {
+    inline HRESULT checked(HRESULT result, const std::source_location loc = std::source_location::current()) {
+        if (result == 0) {
+            return result;
+        }
+
+        auto msg = std::format("Got HRESULT={:#x} at\n{}:{}", static_cast<std::uint32_t>(result) & 0xFFFFFFFF, loc.function_name(), loc.line());
+        throw std::runtime_error(msg);
+    }
+
+    template <typename Callable>
+    inline HRESULT retry_while_pending(Callable&& fn) {
+        bool delayed = false;
+        HRESULT status = 0;
+        do {
+            if (status != 0) {
+                delayed = true;
+                std::println("delaying for com retry...");
+                std::this_thread::sleep_for(std::chrono::milliseconds(500));
+            }
+
+            status = fn();
+        } while (status == E_PENDING);
+
+        return status;
+    }
+
+    /// A very basic implementation, a lot of stuff is missing
+    template <typename Ty>
+    class Ptr {
+    public:
+        Ptr() = default;
+        explicit Ptr(Ty* ptr): ptr_(ptr) { }
+
+        ~Ptr() {
+            release();
+        }
+
+        /// No copying
+        Ptr(const Ptr&) = delete;
+        Ptr& operator=(const Ptr&) = delete;
+
+        /// Move
+        Ptr(Ptr&& other) noexcept: ptr_(other.ptr_) {
+            other.ptr_ = nullptr;
+        }
+        Ptr& operator=(Ptr&& other) noexcept {
+            if (this != &other) {
+                release();
+                ptr_ = other.ptr_;
+                other.ptr_ = nullptr;
+            }
+            return *this;
+        }
+
+    public:
+        [[nodiscard]] Ty* get() const {
+            return ptr_;
+        }
+
+        [[nodiscard]] Ty* operator->() const {
+            return ptr_;
+        }
+
+        [[nodiscard]] Ty** ref_to_ptr() {
+            return &ptr_;
+        }
+
+    private:
+        void release() {
+            if (ptr_ != nullptr) {
+                ptr_->Release();
+            }
+        }
+
+        Ty* ptr_ = nullptr;
+    };
+
+    template <GUID ClsId, GUID IID>
+    class IBaseObject {
+    public:
+        static constexpr GUID kClsId = ClsId;
+        static constexpr GUID kIID = IID;
+
+    private:
+        virtual HRESULT QueryInterface() = 0;
+        virtual std::uint32_t AddRef() = 0;
+
+    public:
+        virtual std::uint32_t Release() = 0;
+    };
+
+    template <typename Ty>
+    concept ComObject = requires {
+        Ty::kClsId;
+        Ty::kIID;
+    };
+
+    template <ComObject Ty>
+    [[nodiscard]] Ptr<Ty> query() {
+        Ptr<Ty> result;
+        const auto status = CoCreateInstance(Ty::kClsId, 0, 1, Ty::kIID, reinterpret_cast<LPVOID*>(result.ref_to_ptr()));
+        if (status == REGDB_E_CLASSNOTREG) {
+            throw std::runtime_error(strings::wsc_unavailable_error().data());
+        }
+
+        checked(status);
+        return result;
+    }
+} // namespace com

defendnot-loader/core/autorun.cpp

@@ -1,5 +1,6 @@
 #include "core/core.hpp"
 
+#include "shared/com.hpp"
 #include "shared/ctx.hpp"
 #include "shared/strings.hpp"
 
@@ -19,38 +20,6 @@ namespace loader {
     namespace {
         constexpr std::string_view kTaskName = strings::kProjectName;
 
-        /// A very basic implementation, a lot of stuff is missing
-        template <typename Ty>
-        class ComPtr {
-        public:
-            ComPtr() = default;
-            explicit ComPtr(Ty* ptr): ptr_(ptr) { }
-
-            ~ComPtr() {
-                if (ptr_ != nullptr) {
-                    ptr_->Release();
-                }
-            }
-
-            ComPtr(const ComPtr&) = delete;
-            ComPtr& operator=(const ComPtr&) = delete;
-
-            [[nodiscard]] Ty* get() const {
-                return ptr_;
-            }
-
-            [[nodiscard]] Ty* operator->() const {
-                return ptr_;
-            }
-
-            [[nodiscard]] Ty** ref_to_ptr() {
-                return &ptr_;
-            }
-
-        private:
-            Ty* ptr_ = nullptr;
-        };
-
         void co_initialize() {
             static std::once_flag fl;
             std::call_once(fl, []() -> void {
@@ -66,7 +35,7 @@ namespace loader {
         [[nodiscard]] bool with_service(Callable&& callback) {
             co_initialize();
 
-            ComPtr<ITaskService> service;
+            com::Ptr<ITaskService> service;
             auto hr =
                 CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast<void**>(service.ref_to_ptr()));
             if (FAILED(hr)) {
@@ -78,7 +47,7 @@ namespace loader {
                 return false;
             }
 
-            ComPtr<ITaskFolder> root_folder;
+            com::Ptr<ITaskFolder> root_folder;
             hr = service->GetFolder(BSTR(L"\\"), root_folder.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
@@ -104,55 +73,55 @@ namespace loader {
                 user_id = bstr_sys;
             }
 
-            ComPtr<ITaskDefinition> task;
+            com::Ptr<ITaskDefinition> task;
             auto hr = service->NewTask(0, task.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<IRegistrationInfo> reg_info;
+            com::Ptr<IRegistrationInfo> reg_info;
             hr = task->get_RegistrationInfo(reg_info.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<IPrincipal> principal;
+            com::Ptr<IPrincipal> principal;
             hr = task->get_Principal(principal.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<ITriggerCollection> trigger_collection;
+            com::Ptr<ITriggerCollection> trigger_collection;
             hr = task->get_Triggers(trigger_collection.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<ITrigger> trigger;
+            com::Ptr<ITrigger> trigger;
             hr = trigger_collection->Create(task_trigger, trigger.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<IActionCollection> action_collection;
+            com::Ptr<IActionCollection> action_collection;
             hr = task->get_Actions(action_collection.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<IAction> action;
+            com::Ptr<IAction> action;
             hr = action_collection->Create(TASK_ACTION_EXEC, action.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<IExecAction> exec_action;
+            com::Ptr<IExecAction> exec_action;
             hr = action->QueryInterface(IID_IExecAction, reinterpret_cast<void**>(exec_action.ref_to_ptr()));
             if (FAILED(hr)) {
                 return false;
             }
 
-            ComPtr<ITaskSettings> settings;
+            com::Ptr<ITaskSettings> settings;
             hr = task->get_Settings(settings.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
@@ -175,7 +144,7 @@ namespace loader {
             exec_action->put_Arguments(bstr_t("--from-autorun"));
 
             /// Register the task and we are done
-            ComPtr<IRegisteredTask> registered_task;
+            com::Ptr<IRegisteredTask> registered_task;
             hr = folder->RegisterTaskDefinition(bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, TASK_LOGON_NONE,
                                                 variant_t(L""), registered_task.ref_to_ptr());
             return SUCCEEDED(hr);

defendnot/bootstrap/bootstrap.cpp

@@ -7,19 +7,35 @@
 #include <Windows.h>
 
 namespace defendnot {
+    namespace {
+        template <com::ComObject Ty>
+        void apply(const std::string_view log_prefix, const BSTR name) {
+            /// Get the WSC interface
+            auto inst = com::query<Ty>();
+
+            /// This can fail if we dont have any products registered so no com_checked
+            logln("{}_unregister: {:#x}", log_prefix, com::retry_while_pending([&inst]() -> HRESULT { return inst->Unregister(); }) & 0xFFFFFFFF);
+            if (shared::ctx.state == shared::State::OFF) {
+                return;
+            }
+
+            /// Register and activate
+            logln("{}_register: {:#x}", log_prefix, com::checked(inst->Register(name, name, 0, 0)));
+            logln("{}_update: {:#x}", log_prefix, com::checked(inst->UpdateStatus(WSCSecurityProductState::ON, static_cast<BOOL>(true))));
+
+            /// Update the substatuses, if the interface supports this
+            if constexpr (std::is_same_v<Ty, IWscAVStatus4>) {
+                logln("{}_scan_update: {:#x}", log_prefix, com::checked(inst->UpdateScanSubstatus(WSCSecurityProductSubStatus::NO_ACTION)));
+                logln("{}_settings_update: {:#x}", log_prefix, com::checked(inst->UpdateSettingsSubstatus(WSCSecurityProductSubStatus::NO_ACTION)));
+                logln("{}_prot_update: {:#x}", log_prefix, com::checked(inst->UpdateProtectionUpdateSubstatus(WSCSecurityProductSubStatus::NO_ACTION)));
+            }
+        }
+    } // namespace
+
     void startup() {
         /// Setup
         shared::ctx.deserialize();
-        logln("init: {:#x}", com_checked(CoInitialize(nullptr)));
+        logln("init: {:#x}", com::checked(CoInitialize(nullptr)));
-
-        /// Get the main WSC interface we will be dealing with
-        auto inst = IWscAVStatus::get();
-
-        /// This can fail if we dont have any avs registered so no com_checked
-        logln("unregister: {:#x}", com_retry_while_pending([&inst]() -> HRESULT { return inst->Unregister(); }) & 0xFFFFFFFF);
-        if (shared::ctx.state == shared::State::OFF) {
-            return;
-        }
 
         /// WSC will reject the register request if name is empty
         auto name_w = std::wstring(shared::ctx.name.begin(), shared::ctx.name.end());
@@ -33,8 +49,8 @@ namespace defendnot {
             SysFreeString(name);
         };
 
-        /// Register and activate our AV
+        /// Register our stuff in the WSC interfaces
-        logln("register: {:#x}", com_checked(inst->Register(name, name, 0, 0)));
+        apply<IWscASStatus>("IWscASStatus", name);
-        logln("update: {:#x}", com_checked(inst->UpdateStatus(WSCSecurityProductState::ON, 3)));
+        apply<IWscAVStatus4>("IWscAVStatus4", name);
     }
 } // namespace defendnot

defendnot/core/com.hpp

@@ -5,14 +5,23 @@
 #include <thread>
 
 #include "core/log.hpp"
+#include "shared/com.hpp"
 #include "shared/strings.hpp"
 
 #include <Windows.h>
 
 namespace defendnot {
     namespace detail {
-        inline GUID CLSID_IWscAVStatus = {0x0F2102C37, 0x90C3, 0x450C, {0x0B3, 0x0F6, 0x92, 0x0BE, 0x16, 0x93, 0x0BD, 0x0F2}};
+        constexpr GUID CLSID_WscIsv = {0xF2102C37, 0x90C3, 0x450C, {0xB3, 0x0F6, 0x92, 0xBE, 0x16, 0x93, 0xBD, 0xF2}};
-        inline GUID IID_IWscAVStatus = {0x3901A765, 0x0AB91, 0x4BA9, {0xA5, 0x53, 0x5B, 0x85, 0x38, 0xDE, 0xB8, 0x40}};
+
+        constexpr GUID IID_IWscFWStatus = {0x9B8F6C6E, 0x8A4A, 0x4891, {0xAF, 0x63, 0x1A, 0x2F, 0x50, 0x92, 0x40, 0x40}};
+        constexpr GUID IID_IWscFWStatus2 = {0x62F698CB, 0x94A, 0x4C68, {0x94, 0x19, 0x8E, 0x8C, 0x49, 0x42, 0x0E, 0x59}};
+
+        constexpr GUID IID_IWscAVStatus = {0x3901A765, 0xAB91, 0x4BA9, {0xA5, 0x53, 0x5B, 0x85, 0x38, 0xDE, 0xB8, 0x40}};
+        constexpr GUID IID_IWscAVStatus3 = {0xCF007CA2, 0xF5E3, 0x11E5, {0x9C, 0xE9, 0x5E, 0x55, 0x17, 0x50, 0x7C, 0x66}};
+        constexpr GUID IID_IWscAVStatus4 = {0x4DCBAFAC, 0x29BA, 0x46B1, {0x80, 0xFC, 0xB8, 0xBD, 0xE3, 0xC0, 0xAE, 0x4D}};
+
+        constexpr GUID IID_IWscASStatus = {0x24E9756, 0xBA6C, 0x4AD1, {0x83, 0x21, 0x87, 0xBA, 0xE7, 0x8F, 0xD0, 0xE3}};
     } // namespace detail
 
     enum class WSCSecurityProductState : std::uint32_t {
@@ -26,49 +35,14 @@ namespace defendnot {
         NOT_SET = 0,
         NO_ACTION = 1,
         ACTION_RECOMMENDED = 2,
-        ACTION_NEEDED = 3
+        ACTION_NEEDED = 3,
     };
 
-    inline HRESULT com_checked(HRESULT result, const std::source_location loc = std::source_location::current()) {
+    class IWscAVStatus4 : public com::IBaseObject<detail::CLSID_WscIsv, detail::IID_IWscAVStatus4> {
-        if (result == 0) {
-            return result;
-        }
-
-        auto msg = std::format("Got HRESULT={:#x} at\n{}:{}", static_cast<std::uint32_t>(result) & 0xFFFFFFFF, loc.function_name(), loc.line());
-        throw std::runtime_error(msg);
-    }
-
-    template <typename Callable>
-    inline HRESULT com_retry_while_pending(Callable&& fn) {
-        bool delayed = false;
-        HRESULT status = 0;
-        do {
-            if (status != 0) {
-                delayed = true;
-                logln("delaying for com retry...");
-                std::this_thread::sleep_for(std::chrono::seconds(5));
-            }
-
-            status = fn();
-        } while (status == E_PENDING);
-
-        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));
-        }
-
-        return status;
-    }
-
-    class IWscAVStatus {
     public:
-        virtual HRESULT QueryInterface() = 0;
-        virtual std::uint32_t AddRef() = 0;
-        virtual std::uint32_t Release() = 0;
         virtual HRESULT Register(BSTR path_to_signed_product_exe, BSTR display_name, std::uint32_t, std::uint32_t) = 0;
         virtual HRESULT Unregister() = 0;
-        virtual HRESULT UpdateStatus(WSCSecurityProductState state, std::uint32_t) = 0;
+        virtual HRESULT UpdateStatus(WSCSecurityProductState state, BOOL unk) = 0;
         virtual HRESULT InitiateOfflineCleaning(std::uint16_t*, std::uint16_t*) = 0;
         virtual HRESULT NotifyUserForNearExpiration(std::uint32_t) = 0;
         virtual HRESULT MakeDefaultProductRequest() = 0;
@@ -78,7 +52,7 @@ namespace defendnot {
         virtual HRESULT UpdateProtectionUpdateSubstatus(WSCSecurityProductSubStatus status) = 0;
         virtual HRESULT RegisterAV(std::uint16_t*, std::uint16_t*, std::uint32_t, std::uint32_t) = 0;
         virtual HRESULT UnregisterAV() = 0;
-        virtual HRESULT UpdateStatusAV(WSCSecurityProductState state, std::uint32_t) = 0;
+        virtual HRESULT UpdateStatusAV(WSCSecurityProductState state, BOOL unk) = 0;
         virtual HRESULT InitiateOfflineCleaningAV(std::uint16_t*, std::uint16_t*) = 0;
         virtual HRESULT NotifyUserForNearExpirationAV(std::uint32_t) = 0;
         virtual HRESULT RegisterFW(std::uint16_t*, std::uint16_t*, std::uint32_t, std::uint32_t) = 0;
@@ -86,21 +60,13 @@ namespace defendnot {
         virtual HRESULT UpdateStatusFW(WSCSecurityProductState state) = 0;
         virtual HRESULT RegisterAS(std::uint16_t*, std::uint16_t*, std::uint32_t, std::uint32_t) = 0;
         virtual HRESULT UnregisterAS() = 0;
-        virtual HRESULT UpdateStatusAS(WSCSecurityProductState state, std::uint32_t) = 0;
+        virtual HRESULT UpdateStatusAS(WSCSecurityProductState state, BOOL unk) = 0;
-
+    };
-    private:
-        virtual void dtor() = 0;
 
+    class IWscASStatus : public com::IBaseObject<detail::CLSID_WscIsv, detail::IID_IWscASStatus> {
     public:
-        static IWscAVStatus* get() {
+        virtual HRESULT Register(BSTR path_to_signed_product_exe, BSTR display_name, std::uint32_t, std::uint32_t) = 0;
-            IWscAVStatus* result = nullptr;
+        virtual HRESULT Unregister() = 0;
-            const auto status = CoCreateInstance(detail::CLSID_IWscAVStatus, 0, 1, detail::IID_IWscAVStatus, reinterpret_cast<LPVOID*>(&result));
+        virtual HRESULT UpdateStatus(WSCSecurityProductState state, BOOL unk) = 0;
-            if (status == REGDB_E_CLASSNOTREG) {
-                throw std::runtime_error(strings::wsc_unavailable_error().data());
-            }
-
-            com_checked(status);
-            return result;
-        }
     };
 } // namespace defendnot