diff --git a/cxx-shared/cxx-shared.vcxitems b/cxx-shared/cxx-shared.vcxitems index 1e2f684..2d57efa 100644 --- a/cxx-shared/cxx-shared.vcxitems +++ b/cxx-shared/cxx-shared.vcxitems @@ -14,6 +14,7 @@ + diff --git a/cxx-shared/shared/com.hpp b/cxx-shared/shared/com.hpp new file mode 100644 index 0000000..65f14d0 --- /dev/null +++ b/cxx-shared/shared/com.hpp @@ -0,0 +1,121 @@ +#pragma once +#include + +#include "shared/strings.hpp" + +#include +#include +#include +#include +#include + +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(result) & 0xFFFFFFFF, loc.function_name(), loc.line()); + throw std::runtime_error(msg); + } + + template + 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 + 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 + 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 + concept ComObject = requires { + Ty::kClsId; + Ty::kIID; + }; + + template + [[nodiscard]] Ptr query() { + Ptr result; + const auto status = CoCreateInstance(Ty::kClsId, 0, 1, Ty::kIID, reinterpret_cast(result.ref_to_ptr())); + if (status == REGDB_E_CLASSNOTREG) { + throw std::runtime_error(strings::wsc_unavailable_error().data()); + } + + checked(status); + return result; + } +} // namespace com diff --git a/defendnot-loader/core/autorun.cpp b/defendnot-loader/core/autorun.cpp index 3d07d2d..4b7db24 100644 --- a/defendnot-loader/core/autorun.cpp +++ b/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 - 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 service; + com::Ptr service; auto hr = CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast(service.ref_to_ptr())); if (FAILED(hr)) { @@ -78,7 +47,7 @@ namespace loader { return false; } - ComPtr root_folder; + com::Ptr 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 task; + com::Ptr task; auto hr = service->NewTask(0, task.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr reg_info; + com::Ptr reg_info; hr = task->get_RegistrationInfo(reg_info.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr principal; + com::Ptr principal; hr = task->get_Principal(principal.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr trigger_collection; + com::Ptr trigger_collection; hr = task->get_Triggers(trigger_collection.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr trigger; + com::Ptr trigger; hr = trigger_collection->Create(task_trigger, trigger.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr action_collection; + com::Ptr action_collection; hr = task->get_Actions(action_collection.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr action; + com::Ptr action; hr = action_collection->Create(TASK_ACTION_EXEC, action.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr exec_action; + com::Ptr exec_action; hr = action->QueryInterface(IID_IExecAction, reinterpret_cast(exec_action.ref_to_ptr())); if (FAILED(hr)) { return false; } - ComPtr settings; + com::Ptr 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 registered_task; + com::Ptr 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); diff --git a/defendnot/bootstrap/bootstrap.cpp b/defendnot/bootstrap/bootstrap.cpp index d869caa..096f3ee 100644 --- a/defendnot/bootstrap/bootstrap.cpp +++ b/defendnot/bootstrap/bootstrap.cpp @@ -7,19 +7,35 @@ #include namespace defendnot { + namespace { + template + void apply(const std::string_view log_prefix, const BSTR name) { + /// Get the WSC interface + auto inst = com::query(); + + /// 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(true)))); + + /// Update the substatuses, if the interface supports this + if constexpr (std::is_same_v) { + 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))); - - /// 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; - } + logln("init: {:#x}", com::checked(CoInitialize(nullptr))); /// 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 - logln("register: {:#x}", com_checked(inst->Register(name, name, 0, 0))); - logln("update: {:#x}", com_checked(inst->UpdateStatus(WSCSecurityProductState::ON, 3))); + /// Register our stuff in the WSC interfaces + apply("IWscASStatus", name); + apply("IWscAVStatus4", name); } } // namespace defendnot diff --git a/defendnot/core/com.hpp b/defendnot/core/com.hpp index b6c84ef..e98fd28 100644 --- a/defendnot/core/com.hpp +++ b/defendnot/core/com.hpp @@ -5,14 +5,23 @@ #include #include "core/log.hpp" +#include "shared/com.hpp" #include "shared/strings.hpp" #include namespace defendnot { namespace detail { - inline GUID CLSID_IWscAVStatus = {0x0F2102C37, 0x90C3, 0x450C, {0x0B3, 0x0F6, 0x92, 0x0BE, 0x16, 0x93, 0x0BD, 0x0F2}}; - inline GUID IID_IWscAVStatus = {0x3901A765, 0x0AB91, 0x4BA9, {0xA5, 0x53, 0x5B, 0x85, 0x38, 0xDE, 0xB8, 0x40}}; + constexpr GUID CLSID_WscIsv = {0xF2102C37, 0x90C3, 0x450C, {0xB3, 0x0F6, 0x92, 0xBE, 0x16, 0x93, 0xBD, 0xF2}}; + + 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()) { - if (result == 0) { - return result; - } - - auto msg = std::format("Got HRESULT={:#x} at\n{}:{}", static_cast(result) & 0xFFFFFFFF, loc.function_name(), loc.line()); - throw std::runtime_error(msg); - } - - template - 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 { + class IWscAVStatus4 : public com::IBaseObject { 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; - - private: - virtual void dtor() = 0; + virtual HRESULT UpdateStatusAS(WSCSecurityProductState state, BOOL unk) = 0; + }; + class IWscASStatus : public com::IBaseObject { public: - static IWscAVStatus* get() { - IWscAVStatus* result = nullptr; - const auto status = CoCreateInstance(detail::CLSID_IWscAVStatus, 0, 1, detail::IID_IWscAVStatus, reinterpret_cast(&result)); - if (status == REGDB_E_CLASSNOTREG) { - throw std::runtime_error(strings::wsc_unavailable_error().data()); - } - - com_checked(status); - return result; - } + 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, BOOL unk) = 0; }; } // namespace defendnot