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}