Support logging actions for interesting modules (#133)

src/analyzer/main.cpp

@@ -15,23 +15,28 @@ namespace
         bool silent{false};
         std::string registry_path{"./registry"};
         std::string emulation_root{};
+        std::set<std::string, std::less<>> modules{};
     };
 
-    void watch_system_objects(windows_emulator& win_emu, const bool cache_logging)
+    void watch_system_objects(windows_emulator& win_emu, const std::set<std::string, std::less<>>& modules,
+                              const bool cache_logging)
     {
         (void)win_emu;
+        (void)modules;
         (void)cache_logging;
 
 #ifdef OS_WINDOWS
-        watch_object(win_emu, *win_emu.current_thread().teb, cache_logging);
-        watch_object(win_emu, win_emu.process().peb, cache_logging);
-        watch_object(win_emu, emulator_object<KUSER_SHARED_DATA64>{win_emu.emu(), kusd_mmio::address()}, cache_logging);
+        watch_object(win_emu, modules, *win_emu.current_thread().teb, cache_logging);
+        watch_object(win_emu, modules, win_emu.process().peb, cache_logging);
+        watch_object(win_emu, modules, emulator_object<KUSER_SHARED_DATA64>{win_emu.emu(), kusd_mmio::address()},
+                     cache_logging);
 
-        auto* params_hook = watch_object(win_emu, win_emu.process().process_params, cache_logging);
+        auto* params_hook = watch_object(win_emu, modules, win_emu.process().process_params, cache_logging);
 
         win_emu.emu().hook_memory_write(
             win_emu.process().peb.value() + offsetof(PEB64, ProcessParameters), 0x8,
-            [&win_emu, cache_logging, params_hook](const uint64_t address, size_t, const uint64_t value) mutable {
+            [&win_emu, cache_logging, params_hook, modules](const uint64_t address, size_t,
+                                                            const uint64_t value) mutable {
                 const auto target_address = win_emu.process().peb.value() + offsetof(PEB64, ProcessParameters);
 
                 if (address == target_address)
@@ -39,7 +44,7 @@ namespace
                     const emulator_object<RTL_USER_PROCESS_PARAMETERS64> obj{win_emu.emu(), value};
 
                     win_emu.emu().delete_hook(params_hook);
-                    params_hook = watch_object(win_emu, obj, cache_logging);
+                    params_hook = watch_object(win_emu, modules, obj, cache_logging);
                 }
             });
 #endif
@@ -108,7 +113,7 @@ namespace
             return false;
         }
 
-        emulator_settings settings{
+        const emulator_settings settings{
             .application = args[0],
             .registry_directory = options.registry_path,
             .emulation_root = options.emulation_root,
@@ -116,12 +121,13 @@ namespace
             .verbose_calls = options.verbose_logging,
             .disable_logging = options.silent,
             .silent_until_main = options.concise_logging,
+            .modules = options.modules,
         };
 
-        windows_emulator win_emu{std::move(settings)};
+        windows_emulator win_emu{settings};
 
         (void)&watch_system_objects;
-        watch_system_objects(win_emu, options.concise_logging);
+        watch_system_objects(win_emu, options.modules, options.concise_logging);
         win_emu.buffer_stdout = true;
 
         if (options.silent)
@@ -226,6 +232,16 @@ namespace
             {
                 options.concise_logging = true;
             }
+            else if (arg == "-m")
+            {
+                if (args.size() < 2)
+                {
+                    throw std::runtime_error("No module provided after -m");
+                }
+
+                arg_it = args.erase(arg_it);
+                options.modules.insert(std::string(args[0]));
+            }
             else if (arg == "-e")
             {
                 if (args.size() < 2)

src/analyzer/object_watching.hpp

@@ -1,18 +1,20 @@
 #pragma once
 
 #include "reflect_type_info.hpp"
+#include <set>
 
 template <typename T>
-emulator_hook* watch_object(windows_emulator& emu, emulator_object<T> object, const bool cache_logging = false)
+emulator_hook* watch_object(windows_emulator& emu, const std::set<std::string, std::less<>>& modules,
+                            emulator_object<T> object, const bool cache_logging = false)
 {
     const reflect_type_info<T> info{};
 
     return emu.emu().hook_memory_read(
         object.value(), object.size(),
-        [i = std::move(info), object, &emu, cache_logging](const uint64_t address, size_t, uint64_t) {
+        [i = std::move(info), object, &emu, cache_logging, modules](const uint64_t address, size_t, uint64_t) {
             const auto rip = emu.emu().read_instruction_pointer();
             const auto* mod = emu.process().mod_manager.find_by_address(rip);
-            const auto is_main_access = mod == emu.process().executable;
+            const auto is_main_access = mod == emu.process().executable || modules.contains(mod->name);
 
             if (!emu.verbose_calls && !is_main_access)
             {

src/common/utils/lazy_object.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include <optional>
+#include <type_traits>
+
+namespace utils
+{
+    template <typename F, typename T>
+    concept CallableWithReturn = requires(const F f) {
+        { f() } -> std::same_as<T>;
+    };
+
+    template <typename T, typename F>
+        requires(CallableWithReturn<F, T>)
+    class lazy_object
+    {
+      public:
+        lazy_object(F accessor)
+            : accessor_(std::move(accessor))
+        {
+        }
+
+        operator const T&() const
+        {
+            return this->get();
+        }
+
+        operator T&()
+        {
+            return this->get();
+        }
+
+        T& operator->()
+            requires std::is_pointer_v<T>
+        {
+            return this->get();
+        }
+
+        const T& operator->() const
+            requires std::is_pointer_v<T>
+        {
+            return this->get();
+        }
+
+      private:
+        F accessor_{};
+        mutable std::optional<T> object_{};
+
+        T& get() const
+        {
+            this->ensure_construction();
+            return *this->object_;
+        }
+
+        void ensure_construction() const
+        {
+            if (!this->object_.has_value())
+            {
+                this->object_.emplace(this->accessor_());
+            }
+        }
+    };
+
+    template <typename F>
+    auto make_lazy(F accessor)
+    {
+        return lazy_object<std::invoke_result_t<F>, F>(std::move(accessor));
+    }
+}

src/windows-emulator/windows_emulator.cpp

@@ -5,9 +5,11 @@
 #include "context_frame.hpp"
 
 #include <unicorn_x64_emulator.hpp>
+
+#include <utils/io.hpp>
 #include <utils/finally.hpp>
-#include "utils/compression.hpp"
-#include "utils/io.hpp"
+#include <utils/compression.hpp>
+#include <utils/lazy_object.hpp>
 
 #include "apiset.hpp"
 
@@ -861,6 +863,8 @@ windows_emulator::windows_emulator(const emulator_settings& settings, emulator_c
     this->use_relative_time_ = settings.use_relative_time;
     this->log.disable_output(settings.disable_logging || this->silent_until_main_);
     this->callbacks_ = std::move(callbacks);
+    this->modules_ = settings.modules;
+
     this->setup_process(settings, working_dir);
 }
 
@@ -955,15 +959,34 @@ void windows_emulator::on_instruction_execution(const uint64_t address)
     const auto thread_insts = ++thread.executed_instructions;
     if (thread_insts % MAX_INSTRUCTIONS_PER_TIME_SLICE == 0)
     {
-        this->switch_thread_ = true;
-        this->emu().stop();
+        this->yield_thread();
     }
 
     process.previous_ip = process.current_ip;
     process.current_ip = this->emu().read_instruction_pointer();
 
+    const auto binary = utils::make_lazy([&] {
+        return this->process().mod_manager.find_by_address(address); //
+    });
+
+    const auto previous_binary = utils::make_lazy([&] {
+        return this->process().mod_manager.find_by_address(process.previous_ip); //
+    });
+
+    const auto is_in_interesting_module = [&] {
+        if (this->modules_.empty())
+        {
+            return false;
+        }
+
+        return (binary && this->modules_.contains(binary->name)) ||
+               (previous_binary && this->modules_.contains(previous_binary->name));
+    };
+
     const auto is_main_exe = process.executable->is_within(address);
-    const auto is_interesting_call = process.executable->is_within(process.previous_ip) || is_main_exe;
+    const auto is_interesting_call = process.executable->is_within(process.previous_ip) //
+                                     || is_main_exe                                     //
+                                     || is_in_interesting_module();
 
     if (this->silent_until_main_ && is_main_exe)
     {
@@ -976,8 +999,6 @@ void windows_emulator::on_instruction_execution(const uint64_t address)
         return;
     }
 
-    const auto* binary = this->process().mod_manager.find_by_address(address);
-
     if (binary)
     {
         const auto export_entry = binary->address_names.find(address);

src/windows-emulator/windows_emulator.hpp

@@ -38,6 +38,7 @@ struct emulator_settings
     bool use_relative_time{false};
     std::unordered_map<uint16_t, uint16_t> port_mappings{};
     std::unordered_map<windows_path, std::filesystem::path> path_mappings{};
+    std::set<std::string, std::less<>> modules{};
 };
 
 enum class apiset_location : uint8_t
@@ -204,6 +205,8 @@ class windows_emulator
     std::vector<instruction_hook_callback> syscall_hooks_{};
     std::unordered_map<uint16_t, uint16_t> port_mappings_{};
 
+    std::set<std::string, std::less<>> modules_{};
+
     process_context process_;
     syscall_dispatcher dispatcher_;