Skip to main content

ib_shell/hook/
inject.rs

1use std::{
2    path::PathBuf,
3    sync::{Arc, Mutex, MutexGuard},
4};
5
6use bon::bon;
7use ib_hook::{
8    inject::dll::app::{DllApp, DllInjectionVecWithInput, OwnedProcess},
9    process::{GuiProcessWatcher, Pid},
10};
11use serde::{Deserialize, Serialize};
12use tracing::error;
13
14use crate::{app::ShellApp, hook::HookConfig};
15
16#[derive(Serialize, Deserialize, Clone, Debug)]
17pub struct Input {
18    pub injector: Pid,
19    pub config: HookConfig,
20}
21
22impl From<HookConfig> for Input {
23    fn from(config: HookConfig) -> Self {
24        Self {
25            injector: Pid::current(),
26            config,
27        }
28    }
29}
30
31pub struct ShellDll;
32impl DllApp for ShellDll {
33    const APPLY: &str = "ib_shell_apply";
34    type Input = Input;
35    type Output = ();
36}
37
38pub struct ShellInjector {
39    _watcher: GuiProcessWatcher,
40    injections: Arc<Mutex<DllInjectionVecWithInput<ShellDll>>>,
41}
42
43#[bon]
44impl ShellInjector {
45    #[builder]
46    pub fn new(
47        dll_path: PathBuf,
48        config: Option<HookConfig>,
49        apps: Vec<ShellApp>,
50    ) -> anyhow::Result<Self> {
51        let input = config.map(Into::into);
52        let injections = DllInjectionVecWithInput::<ShellDll>::with_input(dll_path, input)?;
53        let injections = Arc::new(Mutex::new(injections));
54        let watcher = ib_hook::process::GuiProcessWatcher::for_each({
55            let injections = injections.clone();
56            move |pid| {
57                let mut injections = injections.lock().unwrap();
58                if let Err(e) = injections
59                    .inject([OwnedProcess::from_pid(*pid).unwrap()].into_iter())
60                    .on_error(|pid, e| error!(%pid, ?e))
61                    .call()
62                {
63                    error!(?e);
64                }
65            }
66        })
67        .filter_image_path(move |path| {
68            path.and_then(|p| p.file_name())
69                .is_some_and(|n| apps.iter().any(|app| n == app.process_name))
70        })
71        .build()?;
72        Ok(Self {
73            _watcher: watcher,
74            injections,
75        })
76    }
77
78    pub fn injections<'a>(&'a self) -> MutexGuard<'a, DllInjectionVecWithInput<ShellDll>> {
79        self.injections.lock().unwrap()
80    }
81
82    pub fn apply(&self, config: HookConfig) {
83        self.injections()
84            .apply(Input {
85                injector: Pid::current(),
86                config,
87            })
88            .on_error(|pid, e| error!(%pid, ?e))
89            .call();
90    }
91
92    pub fn unapply(&self) {
93        self.injections()
94            .unapply()
95            .on_error(|pid, e| error!(%pid, ?e))
96            .call();
97    }
98
99    pub fn leak(&self) {
100        self.injections().leak();
101    }
102
103    pub fn eject(&self) {
104        self.injections()
105            .eject()
106            .on_error(|pid, e| error!(%pid, ?e))
107            .call()
108    }
109}