Skip to main content

ib_shell_item/hook/
inject.rs

1use std::{mem::transmute, path::Path};
2
3use anyhow::{Context, anyhow, bail};
4use bon::bon;
5use dll_syringe::{
6    Syringe,
7    process::{BorrowedProcessModule, OwnedProcess},
8    rpc::RemotePayloadProcedure,
9};
10use tracing::{error, info};
11
12use crate::hook::{HookConfig, dll::SET_HOOK};
13
14/// Represents an injected hook with its syringe, payload, and remote set_hook function.
15pub struct ShellItemHook {
16    syringe: Syringe,
17    payload: BorrowedProcessModule<'static>,
18    remote_set_hook: RemotePayloadProcedure<fn(Option<HookConfig>)>,
19}
20
21impl ShellItemHook {
22    /// Enable the hook with the given config.
23    pub fn set_hook(&self, config: &Option<HookConfig>) {
24        let _ = self.remote_set_hook.call(config);
25    }
26
27    /// Disable and detach the hook.
28    pub fn disable_hook(&self) {
29        let _ = self.remote_set_hook.call(&None);
30    }
31
32    /// Eject the DLL from the target process.
33    pub fn eject(self) -> Result<(), String> {
34        self.syringe.eject(self.payload).map_err(|e| e.to_string())
35    }
36}
37
38/// A collection of injected hooks that can be ejected together.
39pub struct ShellItemHooks {
40    hooks: Vec<ShellItemHook>,
41}
42
43#[bon]
44impl ShellItemHooks {
45    /// Inject the hook DLL into all processes with the given name.
46    ///
47    /// Before [`ShellItemHooks::eject()`], the DLL file will be locked and can't be deleted.
48    ///
49    /// # Returns
50    /// - `Ok(ShellItemHooks)` - Successfully injected hooks
51    /// - `Err(anyhow::Error)` - Error during injection
52    #[builder]
53    pub fn inject(
54        /// Path to the hook DLL
55        dll_path: &Path,
56        /// Name of the process to inject into.
57        ///
58        /// e.g., "explorer.exe", "dopus.exe", "Totalcmd64.exe"
59        #[builder(default = "explorer.exe")]
60        process_name: &str,
61        config: Option<HookConfig>,
62    ) -> anyhow::Result<Self> {
63        if !dll_path.exists() {
64            bail!("DLL not found at: {:?}", dll_path);
65        }
66
67        // Find all processes with the given name
68        // TODO: explorer.exe: File explorer only?
69        let processes = OwnedProcess::find_all_by_name(process_name);
70        if processes.is_empty() {
71            bail!("Failed to find any {} process", process_name);
72        }
73        info!("Found {} {} processes", processes.len(), process_name);
74
75        // Store injected hooks for later eject
76        let mut hooks: Vec<ShellItemHook> = Vec::new();
77        for (i, target_process) in processes.into_iter().enumerate() {
78            let syringe = Syringe::for_process(target_process);
79
80            info!("[{}] Injecting DLL: {:?}", i, dll_path);
81            match syringe.find_or_inject(&dll_path) {
82                Ok(payload) => {
83                    info!("[{}] Successfully injected hook.dll", i);
84
85                    // Call set_hook to enable the hook
86                    let remote_set_hook =
87                        unsafe { syringe.get_payload_procedure(payload, SET_HOOK) }
88                            .context("Failed to get set_hook procedure")?
89                            .context("set_hook not found")?;
90
91                    let injected_hook = ShellItemHook {
92                        payload: unsafe { transmute(payload) },
93                        syringe,
94                        remote_set_hook,
95                    };
96                    if config.is_some() {
97                        injected_hook.set_hook(&config);
98                        info!("[{}] Hook enabled", i);
99                    }
100
101                    hooks.push(injected_hook);
102                }
103                Err(e) => {
104                    error!("[{}] Failed to inject DLL: {}", i, e);
105                }
106            }
107        }
108
109        Ok(ShellItemHooks { hooks })
110    }
111
112    /// Eject all hooks and return the first error if any.
113    pub fn eject(&mut self) -> Result<(), anyhow::Error> {
114        let mut first_error: Option<anyhow::Error> = None;
115
116        for hook in self.hooks.drain(..) {
117            hook.disable_hook();
118            if let Err(e) = hook.eject() {
119                error!("Failed to eject DLL: {}", e);
120                if first_error.is_none() {
121                    first_error = Some(anyhow::anyhow!(e));
122                }
123            }
124        }
125
126        if let Some(e) = first_error {
127            return Err(anyhow!(e));
128        }
129
130        info!("Successfully ejected DLL");
131        Ok(())
132    }
133}