Skip to content

Commit 0c6fe45

Browse files
radiantlyradiantly
committed
Add icons to application list
1 parent 619ec23 commit 0c6fe45

File tree

5 files changed

+182
-36
lines changed

5 files changed

+182
-36
lines changed

‎injector/src/cli.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::injector;
1+
use crate::native;
22
use clap::{ArgAction, Args, Error, Parser, error::ErrorKind};
33
use dll_syringe::process::{OwnedProcess, Process};
44
use std::collections::HashMap;
@@ -79,7 +79,7 @@ pub fn start() {
7979

8080
// populate available windows
8181
let mut windows: HashMap<u32, Vec<u32>> = HashMap::new();
82-
injector::get_top_level_windows()
82+
native::get_top_level_windows()
8383
.into_iter()
8484
.for_each(|window_info| {
8585
windows
@@ -90,7 +90,7 @@ pub fn start() {
9090

9191
processes.into_iter().for_each(|(pid, process)| {
9292
if let Some(hwnds) = windows.remove(&pid) {
93-
injector::set_window_props(process, &hwnds, cli.hide_args.hide, false);
93+
native::set_window_props(process, &hwnds, cli.hide_args.hide, false);
9494
} else {
9595
print_error(format!(
9696
"Cannot find any top level windows for pid {:?}",

‎injector/src/gui.rs‎

Lines changed: 85 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
1-
use crate::injector::{self, WindowInfo};
1+
use crate::native::{self, WindowInfo};
22
use eframe::{
33
Renderer,
44
egui::{
5-
self, Color32, ColorImage, Direction, FontData, FontDefinitions, FontFamily, FontId,
6-
IconData, Layout, Margin, RichText, TextStyle, Theme,
5+
self, Atom, AtomExt, Color32, ColorImage, Direction, FontData, FontDefinitions, FontFamily,
6+
FontId, IconData, Image, Layout, Margin, RichText, TextStyle, Theme, Vec2,
77
},
88
};
99
use image::{GenericImageView, ImageFormat, ImageReader};
10-
use std::sync::{Arc, Mutex};
1110
use std::thread;
11+
use std::{
12+
collections::HashMap,
13+
sync::{Arc, Mutex},
14+
};
1215
use std::{io::Cursor, mem};
1316
use windows_capture::{
1417
capture::{CaptureControl, Context, GraphicsCaptureApiHandler},
@@ -21,8 +24,8 @@ use windows_capture::{
2124
};
2225

2326
enum CaptureWorkerEvent {
24-
CAPTURE(Monitor),
25-
NONE,
27+
Capture(Monitor),
28+
StopCapture,
2629
}
2730

2831
struct ScreenCapture {
@@ -81,6 +84,7 @@ struct Gui {
8184
hide_from_taskbar: bool,
8285
show_desktop_preview: bool,
8386
active_monitor: usize,
87+
icon_cache: HashMap<(u32, u32), Option<egui::TextureHandle>>,
8488
}
8589

8690
impl Gui {
@@ -95,17 +99,17 @@ impl Gui {
9599
match event {
96100
InjectorWorkerEvent::UPDATE => {
97101
println!("populating");
98-
let mut w = injector::get_top_level_windows();
102+
let mut w = native::get_top_level_windows();
99103
*windows_copy.lock().unwrap() = mem::take(&mut w);
100104
println!("populating done");
101105
}
102106
InjectorWorkerEvent::HIDE(pid, hwnd, show_on_taskbar) => {
103107
println!("wanna hide {:?}", hwnd);
104-
injector::set_window_props_with_pid(pid, hwnd, true, show_on_taskbar);
108+
native::set_window_props_with_pid(pid, hwnd, true, show_on_taskbar);
105109
}
106110
InjectorWorkerEvent::SHOW(pid, hwnd, show_on_taskbar) => {
107111
println!("wanna show {:?}", hwnd);
108-
injector::set_window_props_with_pid(pid, hwnd, false, show_on_taskbar);
112+
native::set_window_props_with_pid(pid, hwnd, false, show_on_taskbar);
109113
}
110114
}
111115
}
@@ -123,7 +127,7 @@ impl Gui {
123127
}
124128

125129
match event {
126-
CaptureWorkerEvent::CAPTURE(monitor) => {
130+
CaptureWorkerEvent::Capture(monitor) => {
127131
let settings = Settings::new(
128132
monitor,
129133
CursorCaptureSettings::Default,
@@ -139,15 +143,15 @@ impl Gui {
139143
active_capture_control = Some(capture_control);
140144
}
141145
}
142-
CaptureWorkerEvent::NONE => (),
146+
CaptureWorkerEvent::StopCapture => (),
143147
}
144148
}
145149
});
146150

147151
let monitors = Monitor::enumerate().unwrap_or_default();
148152

149153
if monitors.len() > 0 {
150-
let _ = capture_event_send.send(CaptureWorkerEvent::CAPTURE(monitors[0]));
154+
let _ = capture_event_send.send(CaptureWorkerEvent::Capture(monitors[0]));
151155
}
152156

153157
Gui {
@@ -160,25 +164,46 @@ impl Gui {
160164
capture_tex: None,
161165
hide_from_taskbar: false,
162166
active_monitor: 0,
167+
icon_cache: HashMap::new(),
163168
}
164169
}
165170

171+
fn get_icon<'a>(
172+
icon_cache: &'a mut HashMap<(u32, u32), Option<egui::TextureHandle>>,
173+
ctx: &egui::Context,
174+
pid: u32,
175+
hwnd: u32,
176+
) -> &'a Option<egui::TextureHandle> {
177+
if !icon_cache.contains_key(&(pid, hwnd)) {
178+
let icon = match native::get_icon(hwnd) {
179+
Some((width, height, buffer)) => {
180+
let image = ColorImage::from_rgba_unmultiplied([width, height], &buffer);
181+
Some(ctx.load_texture("icon", image, egui::TextureOptions::LINEAR))
182+
}
183+
None => None,
184+
};
185+
186+
icon_cache.insert((pid, hwnd), icon);
187+
}
188+
189+
icon_cache.get(&(pid, hwnd)).unwrap()
190+
}
191+
166192
fn add_section_header(
167193
ui: &mut egui::Ui,
168194
theme: Theme,
169195
header: impl Into<String>,
170196
desc: impl Into<String>,
171197
) {
172-
let (header_color, desc_color) = if theme == Theme::Light {
173-
(
198+
let (header_color, desc_color) = match theme {
199+
Theme::Light => (
174200
Color32::from_rgb(34, 34, 34),
175201
Color32::from_rgb(119, 119, 119),
176-
)
177-
} else {
178-
(
202+
),
203+
Theme::Dark => (
179204
Color32::from_rgb(242, 242, 242),
180205
Color32::from_rgb(148, 148, 148),
181-
)
206+
),
182207
};
183208

184209
ui.label(RichText::new(header).heading().color(header_color));
@@ -189,25 +214,36 @@ impl Gui {
189214

190215
impl eframe::App for Gui {
191216
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
192-
let theme = ctx.theme();
193-
let (events, focused) = ctx.input(|i| (i.events.clone(), i.focused));
217+
// graceful close
218+
if ctx.input(|i| i.viewport().close_requested()) {
219+
let _ = self
220+
.capture_event_send
221+
.send(CaptureWorkerEvent::StopCapture);
222+
return;
223+
}
194224

195-
for event in events {
225+
// check if focus has changed
226+
for event in ctx.input(|i| i.events.clone()) {
196227
if let egui::Event::WindowFocused(focused) = event {
197228
if focused {
198229
println!("focused");
199230
self.event_sender.send(InjectorWorkerEvent::UPDATE).unwrap();
200231
if self.show_desktop_preview {
201-
let _ = self.capture_event_send.send(CaptureWorkerEvent::CAPTURE(
232+
let _ = self.capture_event_send.send(CaptureWorkerEvent::Capture(
202233
self.monitors[self.active_monitor],
203234
));
204235
}
205236
} else {
206-
let _ = self.capture_event_send.send(CaptureWorkerEvent::NONE);
237+
let _ = self
238+
.capture_event_send
239+
.send(CaptureWorkerEvent::StopCapture);
207240
}
208241
}
209242
}
210243

244+
let theme = ctx.theme();
245+
let focused = ctx.input(|i| i.focused);
246+
211247
egui::CentralPanel::default()
212248
.frame(egui::Frame::central_panel(&ctx.style()).inner_margin(Margin::same(14)))
213249
.show(ctx, |ui| {
@@ -259,7 +295,7 @@ impl eframe::App for Gui {
259295
self.active_monitor = i;
260296
let _ = self
261297
.capture_event_send
262-
.send(CaptureWorkerEvent::CAPTURE(*monitor));
298+
.send(CaptureWorkerEvent::Capture(*monitor));
263299
}
264300
}
265301
});
@@ -277,8 +313,29 @@ impl eframe::App for Gui {
277313

278314
for window_info in self.windows.lock().unwrap().iter_mut() {
279315
ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Truncate); // elide with "…"
316+
317+
let icon_atom = if let Some(texture) = Gui::get_icon(
318+
&mut self.icon_cache,
319+
ctx,
320+
window_info.pid,
321+
window_info.hwnd,
322+
) {
323+
Image::from_texture(texture)
324+
.max_height(16.0)
325+
.atom_max_width(16.0)
326+
} else {
327+
Atom::grow().atom_size(Vec2::new(16.0, 0.0))
328+
};
329+
330+
let checkbox_label = (
331+
Atom::grow().atom_size(Vec2::new(0.0, 0.0)),
332+
icon_atom,
333+
Atom::grow().atom_size(Vec2::new(0.0, 0.0)),
334+
&window_info.title,
335+
);
336+
280337
let checkbox_response =
281-
ui.checkbox(&mut window_info.hidden, &window_info.title);
338+
ui.checkbox(&mut window_info.hidden, checkbox_label);
282339
if checkbox_response.changed() {
283340
let event = if window_info.hidden {
284341
InjectorWorkerEvent::HIDE(
@@ -295,6 +352,7 @@ impl eframe::App for Gui {
295352
};
296353
self.event_sender.send(event).unwrap();
297354
}
355+
ui.add_space(2.0);
298356
}
299357
ui.add_space(10.0);
300358
ui.collapsing("Advanced settings", |ui| {
@@ -303,10 +361,10 @@ impl eframe::App for Gui {
303361
ui.checkbox(&mut self.show_desktop_preview, "Show desktop preview");
304362
if preview_checkbox_response.changed() {
305363
let event = if self.show_desktop_preview {
306-
CaptureWorkerEvent::CAPTURE(self.monitors[self.active_monitor])
364+
CaptureWorkerEvent::Capture(self.monitors[self.active_monitor])
307365
} else {
308366
self.capture_tex = None;
309-
CaptureWorkerEvent::NONE
367+
CaptureWorkerEvent::StopCapture
310368
};
311369
self.capture_event_send.send(event).unwrap();
312370
}

‎injector/src/main.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
mod cli;
44
mod gui;
5-
mod injector;
5+
mod native;
66

77
use std::env;
88

injector/src/injector.rs renamed to injector/src/native.rs

Lines changed: 92 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,18 @@ use dll_syringe::{
66
use std::env;
77
use windows::{
88
Win32::{
9-
Foundation::{HWND, LPARAM, TRUE},
10-
Graphics::Dwm::{DWMWA_CLOAKED, DwmGetWindowAttribute},
9+
Foundation::{HWND, LPARAM, TRUE, WPARAM},
10+
Graphics::{
11+
Dwm::{DWMWA_CLOAKED, DwmGetWindowAttribute},
12+
Gdi::{
13+
BI_RGB, BITMAP, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, DeleteObject, GetDC,
14+
GetDIBits, GetObjectW, ReleaseDC,
15+
},
16+
},
1117
UI::WindowsAndMessaging::{
12-
EnumWindows, GetWindowDisplayAffinity, GetWindowTextW, GetWindowThreadProcessId,
13-
IsWindowVisible,
18+
EnumWindows, GCLP_HICONSM, GetClassLongPtrW, GetIconInfo, GetWindowDisplayAffinity,
19+
GetWindowTextW, GetWindowThreadProcessId, HICON, ICON_SMALL2, ICONINFO,
20+
IsWindowVisible, SendMessageW, WM_GETICON,
1421
},
1522
},
1623
core::BOOL,
@@ -24,6 +31,86 @@ pub struct WindowInfo {
2431
pub hidden: bool,
2532
}
2633

34+
pub fn get_icon(hwnd: u32) -> Option<(usize, usize, Vec<u8>)> {
35+
let hwnd = HWND(hwnd.clone() as *mut _);
36+
let lresult =
37+
unsafe { SendMessageW(hwnd, WM_GETICON, Some(WPARAM(ICON_SMALL2 as usize)), None) };
38+
39+
let hicon = if lresult.0 == 0 {
40+
println!("- no hicon from sendmessage");
41+
42+
let uresult = unsafe { GetClassLongPtrW(hwnd, GCLP_HICONSM) };
43+
if uresult == 0 {
44+
println!("- no hicon from getclasslongptrsm");
45+
return None;
46+
}
47+
HICON(uresult as *mut _)
48+
} else {
49+
HICON(lresult.0 as *mut _)
50+
};
51+
52+
let mut icon_info = ICONINFO::default();
53+
let info_result = unsafe { GetIconInfo(hicon, &mut icon_info as *mut _) };
54+
if let Err(err) = info_result {
55+
println!("- no iconinfo retrieved {:?}", err);
56+
return None;
57+
}
58+
59+
let hdc = unsafe { GetDC(None) };
60+
if hdc.is_invalid() {
61+
println!("- no dc");
62+
return None;
63+
}
64+
65+
let mut bitmap = BITMAP::default();
66+
let object_result = unsafe {
67+
GetObjectW(
68+
icon_info.hbmColor.into(),
69+
std::mem::size_of::<BITMAP>() as i32,
70+
Some(&mut bitmap as *mut _ as *mut _),
71+
)
72+
};
73+
if object_result == 0 {
74+
println!("no object");
75+
return None;
76+
}
77+
78+
let mut bmi = BITMAPINFO::default();
79+
bmi.bmiHeader.biSize = std::mem::size_of::<BITMAPINFOHEADER>() as u32;
80+
bmi.bmiHeader.biWidth = bitmap.bmWidth;
81+
bmi.bmiHeader.biHeight = -bitmap.bmHeight;
82+
bmi.bmiHeader.biPlanes = 1;
83+
bmi.bmiHeader.biBitCount = 32;
84+
bmi.bmiHeader.biCompression = BI_RGB.0;
85+
86+
let pixel_count = bitmap.bmWidth * bitmap.bmHeight;
87+
let mut pixels: Vec<u8> = vec![0; (pixel_count * 4) as usize];
88+
let _ = unsafe {
89+
GetDIBits(
90+
hdc,
91+
icon_info.hbmColor,
92+
0,
93+
bitmap.bmHeight as u32,
94+
Some(pixels.as_mut_ptr() as *mut _),
95+
&mut bmi as *mut _,
96+
DIB_RGB_COLORS,
97+
)
98+
};
99+
100+
for i in (0..pixels.len()).step_by(4) {
101+
(pixels[i], pixels[i + 1], pixels[i + 2], pixels[i + 3]) =
102+
(pixels[i + 2], pixels[i + 1], pixels[i], pixels[i + 3]);
103+
}
104+
105+
let icon = Some((bitmap.bmWidth as usize, bitmap.bmHeight as usize, pixels));
106+
107+
let _ = unsafe { ReleaseDC(None, hdc) };
108+
let _ = unsafe { DeleteObject(icon_info.hbmColor.into()) };
109+
let _ = unsafe { DeleteObject(icon_info.hbmMask.into()) };
110+
111+
return icon;
112+
}
113+
27114
unsafe extern "system" fn enum_windows_proc(hwnd: HWND, lparam: LPARAM) -> BOOL {
28115
// skip invisible windows
29116
let is_visible = unsafe { IsWindowVisible(hwnd) };
@@ -101,7 +188,7 @@ where
101188
{
102189
let mut dll_path = env::current_exe().unwrap();
103190
dll_path.pop();
104-
dll_path.push("payload.dll");
191+
dll_path.push("utils.dll");
105192

106193
let injected_payload = syringe.find_or_inject(dll_path).unwrap();
107194

0 commit comments

Comments
 (0)