Merge pull request #10 from es3n1n/autorun-changes

refactor(autorun): create scheduled task to be launched as system

cxx-shared/shared/names.hpp

@@ -7,4 +7,6 @@ namespace names {
 
     constexpr std::string_view kVictimProcess = "Taskmgr.exe";
     constexpr std::string_view kDllName = "defendnot.dll";
+
+    constexpr std::string_view kVersion = "1.0.0";
 } // namespace names

defendnot-loader/core/autorun.cpp

@@ -90,10 +90,20 @@ namespace loader {
         }
     } // namespace
 
-    [[nodiscard]] bool add_to_autorun() {
+    [[nodiscard]] bool add_to_autorun(AutorunType type) {
         const auto bin_path = shared::get_this_module_path();
 
-        return with_service([bin_path](ITaskService* service, ITaskFolder* folder) -> bool {
+        return with_service([bin_path, type](ITaskService* service, ITaskFolder* folder) -> bool {
+            /// Deduce the autorun config based on type
+            const auto task_trigger = type == AutorunType::AS_SYSTEM_ON_BOOT ? TASK_TRIGGER_BOOT : TASK_TRIGGER_LOGON;
+            const auto logon_type = type == AutorunType::AS_SYSTEM_ON_BOOT ? TASK_LOGON_SERVICE_ACCOUNT : TASK_LOGON_INTERACTIVE_TOKEN;
+
+            BSTR user_id = nullptr;
+            auto bstr_sys = bstr_t(L"SYSTEM");
+            if (type == AutorunType::AS_SYSTEM_ON_BOOT) {
+                user_id = bstr_sys;
+            }
+
             ComPtr<ITaskDefinition> task;
             auto hr = service->NewTask(0, task.ref_to_ptr());
             if (FAILED(hr)) {
@@ -119,7 +129,7 @@ namespace loader {
             }
 
             ComPtr<ITrigger> trigger;
-            hr = trigger_collection->Create(TASK_TRIGGER_LOGON, trigger.ref_to_ptr());
+            hr = trigger_collection->Create(task_trigger, trigger.ref_to_ptr());
             if (FAILED(hr)) {
                 return false;
             }
@@ -148,16 +158,26 @@ namespace loader {
                 return false;
             }
 
+            /// Elevated, when system boots
+            principal->put_UserId(user_id);
+            principal->put_LogonType(logon_type);
             principal->put_RunLevel(TASK_RUNLEVEL_HIGHEST);
-            reg_info->put_Author(_bstr_t(names::kRepoUrl.data()));
+
+            /// Info
+            reg_info->put_Author(bstr_t(names::kRepoUrl.data()));
+
+            /// Start even if we're on batteries
             settings->put_DisallowStartIfOnBatteries(VARIANT_FALSE);
             settings->put_StopIfGoingOnBatteries(VARIANT_FALSE);
-            exec_action->put_Path(_bstr_t(bin_path.string().c_str()));
-            exec_action->put_Arguments(_bstr_t("--from-autorun"));
 
+            /// Binary
+            exec_action->put_Path(bstr_t(bin_path.string().c_str()));
+            exec_action->put_Arguments(bstr_t("--from-autorun"));
+
+            /// Register the task and we are done
             ComPtr<IRegisteredTask> registered_task;
-            hr = folder->RegisterTaskDefinition(_bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{},
+            hr = folder->RegisterTaskDefinition(bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, TASK_LOGON_NONE,
-                                                TASK_LOGON_INTERACTIVE_TOKEN, _variant_t(L""), registered_task.ref_to_ptr());
+                                                variant_t(L""), registered_task.ref_to_ptr());
             return SUCCEEDED(hr);
         });
     }

defendnot-loader/core/core.hpp

@@ -5,15 +5,22 @@
 #include <Windows.h>
 
 namespace loader {
+    enum class AutorunType : std::uint8_t {
+        AS_SYSTEM_ON_BOOT = 0, ///< launch on system boot as NT AUTHORITY\SYSTEM
+        AS_CURRENT_USER_ON_LOGIN = 1, ///< launch on user login
+    };
+
     struct Config {
     public:
         std::string name;
         bool disable;
         bool verbose;
         bool from_autorun;
+        AutorunType autorun_type;
+        bool enable_autorun;
     };
 
     [[nodiscard]] HANDLE inject(std::string_view dll_path, std::string_view proc_name);
-    [[nodiscard]] bool add_to_autorun();
+    [[nodiscard]] bool add_to_autorun(AutorunType type);
     [[nodiscard]] bool remove_from_autorun();
 } // namespace loader

defendnot-loader/main.cpp

@@ -56,8 +56,8 @@ namespace {
     }
 
     void process_autorun(const loader::Config& config) {
-        if (shared::ctx.state == shared::State::ON) {
+        if (shared::ctx.state == shared::State::ON && config.enable_autorun) {
-            std::println("** added to autorun: {}", loader::add_to_autorun());
+            std::println("** added to autorun: {}", loader::add_to_autorun(config.autorun_type));
         } else {
             std::println("** removed from autorun: {}", loader::remove_from_autorun());
         }
@@ -75,19 +75,42 @@ namespace {
 } // namespace
 
 int main(int argc, char* argv[]) try {
-    argparse::ArgumentParser program(std::format("{}-loader", names::kProjectName), "1.0.0");
+    argparse::ArgumentParser program(std::format("{}-loader", names::kProjectName), names::kVersion.data(), argparse::default_arguments::none);
 
+    const auto fatal_print = [](const std::string_view str) -> void {
+        shared::alloc_console();
+        std::cerr << str << std::endl;
+        system("pause");
+        std::exit(EXIT_FAILURE);
+    };
+
+    /// We are registering these ourselves because we have to alloc console first
+    program.add_argument("-h", "--help")
+        .help("prints help message and exits")
+        .default_value(false)
+        .implicit_value(true)
+        .action([&fatal_print, &program](const auto& /*unused*/) -> void { fatal_print(program.help().str()); });
+    program.add_argument("--version")
+        .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)); });
+
+    /// defendnot-loader parameters:
     program.add_argument("-n", "--name").help("av display name").default_value(std::string(names::kRepoUrl)).nargs(1);
     program.add_argument("-d", "--disable").help(std::format("disable {}", names::kProjectName)).default_value(false).implicit_value(true);
     program.add_argument("-v", "--verbose").help("verbose logging").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);
+    program.add_argument("--disable-autorun").help("disable autorun task creation").default_value(false).implicit_value(true);
     program.add_argument("--from-autorun").hidden().default_value(false).implicit_value(true);
 
     try {
         program.parse_args(argc, argv);
-    } catch (...) {
+    } catch (std::exception& e) {
-        shared::alloc_console();
+        std::stringstream ss;
-        std::cerr << program;
+        ss << e.what() << '\n';
-        system("pause");
+        ss << program.help().str();
+        fatal_print(ss.str());
         return EXIT_FAILURE;
     }
 
@@ -96,6 +119,10 @@ int main(int argc, char* argv[]) try {
         .disable = program.get<bool>("-d"),
         .verbose = program.get<bool>("-v"),
         .from_autorun = program.get<bool>("--from-autorun"),
+        .autorun_type = program.get<bool>("--autorun-as-user") ? /// As system on boot is the default value
+                            loader::AutorunType::AS_CURRENT_USER_ON_LOGIN :
+                            loader::AutorunType::AS_SYSTEM_ON_BOOT,
+        .enable_autorun = !program.get<bool>("--disable-autorun"),
     };
 
     setup_window(config);