From 9cc344f99d1671ae33c026783ed8d9767fe8c5ce Mon Sep 17 00:00:00 2001 From: Louis Beaumont Date: Fri, 15 Nov 2024 11:08:40 -0800 Subject: [PATCH 1/8] feat: rdp integration --- .../src/bin/screenpipe-server.rs | 12 +- screenpipe-server/src/cli.rs | 6 +- screenpipe-server/src/core.rs | 26 ++-- screenpipe-server/src/video.rs | 18 +-- screenpipe-vision/Cargo.toml | 2 + .../src/bin/screenpipe-vision.rs | 25 +++- screenpipe-vision/src/core.rs | 135 ++++++++++++++---- screenpipe-vision/src/lib.rs | 4 +- 8 files changed, 167 insertions(+), 61 deletions(-) diff --git a/screenpipe-server/src/bin/screenpipe-server.rs b/screenpipe-server/src/bin/screenpipe-server.rs index 9c986a0c40..fe5a940c9c 100644 --- a/screenpipe-server/src/bin/screenpipe-server.rs +++ b/screenpipe-server/src/bin/screenpipe-server.rs @@ -15,7 +15,7 @@ use screenpipe_core::find_ffmpeg_path; use screenpipe_server::{ cli::{Cli, CliAudioTranscriptionEngine, CliOcrEngine, Command, PipeCommand}, start_continuous_recording, watch_pid, DatabaseManager, PipeControl, PipeManager, ResourceMonitor, Server, highlight::{Highlight,HighlightConfig} }; -use screenpipe_vision::{monitor::list_monitors}; +use screenpipe_vision::{core::CaptureSource, monitor::{get_default_monitor, list_monitors}}; #[cfg(target_os = "macos")] use screenpipe_vision::run_ui; use serde_json::{json, Value}; @@ -347,7 +347,6 @@ async fn main() -> anyhow::Result<()> { let vision_control_clone = Arc::clone(&vision_control); let shutdown_tx_clone = shutdown_tx.clone(); let friend_wearable_uid_clone: Option = friend_wearable_uid.clone(); // Clone here - let monitor_ids_clone = monitor_ids.clone(); let ignored_windows_clone = cli.ignored_windows.clone(); let included_windows_clone = cli.included_windows.clone(); @@ -366,6 +365,13 @@ async fn main() -> anyhow::Result<()> { loop { let vad_engine_clone = vad_engine.clone(); // Clone it here for each iteration let mut shutdown_rx = shutdown_tx_clone.subscribe(); + let capture_source = if cli.rdp_session_id.len() > 0 { + cli.rdp_session_id.iter().map(|s| CaptureSource::RdpSession(Box::leak(s.clone().into_boxed_str()))).collect() + } else if cli.monitor_id.len() > 0 { + cli.monitor_id.iter().map(|id| CaptureSource::LocalMonitor(*id)).collect() + } else { + vec![CaptureSource::LocalMonitor(get_default_monitor().await.id())] + }; let recording_future = start_continuous_recording( db_clone.clone(), output_path_clone.clone(), @@ -379,7 +385,7 @@ async fn main() -> anyhow::Result<()> { Arc::new(cli.audio_transcription_engine.clone().into()), Arc::new(cli.ocr_engine.clone().into()), friend_wearable_uid_clone.clone(), - monitor_ids_clone.clone(), + capture_source, cli.use_pii_removal, cli.disable_vision, vad_engine_clone, diff --git a/screenpipe-server/src/cli.rs b/screenpipe-server/src/cli.rs index 077021165b..f1b97a83d7 100644 --- a/screenpipe-server/src/cli.rs +++ b/screenpipe-server/src/cli.rs @@ -243,13 +243,15 @@ pub struct Cli { #[arg(long, default_value_t = false)] pub enable_frame_cache: bool, + /// RDP session ID to use for screen capture - this will override the monitor ID + #[arg(long)] + pub rdp_session_id: Vec, + #[command(subcommand)] pub command: Option, - } - impl Cli { pub fn unique_languages(&self) -> Result, String> { let mut unique_langs = std::collections::HashSet::new(); diff --git a/screenpipe-server/src/core.rs b/screenpipe-server/src/core.rs index 5d4474a5cf..824bd43161 100644 --- a/screenpipe-server/src/core.rs +++ b/screenpipe-server/src/core.rs @@ -13,6 +13,7 @@ use screenpipe_audio::{ use screenpipe_core::pii_removal::remove_pii; use screenpipe_core::Language; use screenpipe_integrations::friend_wearable::initialize_friend_wearable_loop; +use screenpipe_vision::core::CaptureSource; use screenpipe_vision::OcrEngine; use std::collections::HashMap; use std::path::PathBuf; @@ -35,7 +36,7 @@ pub async fn start_continuous_recording( audio_transcription_engine: Arc, ocr_engine: Arc, friend_wearable_uid: Option, - monitor_ids: Vec, + capture_sources: Vec>, use_pii_removal: bool, vision_disabled: bool, vad_engine: CliVadEngine, @@ -55,11 +56,11 @@ pub async fn start_continuous_recording( )); } - debug!("Starting video recording for monitor {:?}", monitor_ids); + debug!("Starting video recording for {:?}", capture_sources); let video_tasks = if !vision_disabled { - monitor_ids - .iter() - .map(|&monitor_id| { + capture_sources + .into_iter() + .map(|capture_source| { let db_manager_video = Arc::clone(&db); let output_path_video = Arc::clone(&output_path); let is_running_video = Arc::clone(&vision_control); @@ -70,7 +71,7 @@ pub async fn start_continuous_recording( let languages = languages.clone(); - debug!("Starting video recording for monitor {}", monitor_id); + debug!("Starting video recording for {}", capture_source); vision_handle.spawn(async move { record_video( db_manager_video, @@ -80,7 +81,7 @@ pub async fn start_continuous_recording( save_text_files, ocr_engine, friend_wearable_uid_video, - monitor_id, + capture_source, use_pii_removal, &ignored_windows_video, &include_windows_video, @@ -180,7 +181,7 @@ async fn record_video( save_text_files: bool, ocr_engine: Arc, _friend_wearable_uid: Option, - monitor_id: u32, + capture_source: CaptureSource<'static>, use_pii_removal: bool, ignored_windows: &[String], include_windows: &[String], @@ -190,7 +191,7 @@ async fn record_video( debug!("record_video: Starting"); let db_chunk_callback = Arc::clone(&db); let rt = Handle::current(); - let device_name = Arc::new(format!("monitor_{}", monitor_id)); + let device_name = Arc::new(format!("{}", capture_source)); let new_chunk_callback = { let db_chunk_callback = Arc::clone(&db_chunk_callback); @@ -200,7 +201,10 @@ async fn record_video( let db_chunk_callback = Arc::clone(&db_chunk_callback); let device_name = Arc::clone(&device_name); rt.spawn(async move { - if let Err(e) = db_chunk_callback.insert_video_chunk(&file_path, &device_name).await { + if let Err(e) = db_chunk_callback + .insert_video_chunk(&file_path, &device_name) + .await + { error!("Failed to insert new video chunk: {}", e); } debug!("record_video: Inserted new video chunk: {}", file_path); @@ -215,7 +219,7 @@ async fn record_video( new_chunk_callback, save_text_files, Arc::clone(&ocr_engine), - monitor_id, + capture_source, ignored_windows, include_windows, languages, diff --git a/screenpipe-server/src/video.rs b/screenpipe-server/src/video.rs index 6c155c4761..dec6b5171f 100644 --- a/screenpipe-server/src/video.rs +++ b/screenpipe-server/src/video.rs @@ -4,6 +4,7 @@ use image::ImageFormat::{self}; use log::{debug, error}; use log::{info, warn}; use screenpipe_core::{find_ffmpeg_path, Language}; +use screenpipe_vision::core::CaptureSource; use screenpipe_vision::{continuous_capture, CaptureResult, OcrEngine}; use std::path::PathBuf; use std::process::Stdio; @@ -33,7 +34,7 @@ impl VideoCapture { new_chunk_callback: impl Fn(&str) + Send + Sync + 'static, save_text_files: bool, ocr_engine: Arc, - monitor_id: u32, + capture_source: CaptureSource<'static>, ignore_list: &[String], include_list: &[String], languages: Vec, @@ -61,7 +62,7 @@ impl VideoCapture { interval, save_text_files, *ocr_engine, - monitor_id, + capture_source, &ignore_list_clone, &include_list_clone, languages.clone(), @@ -69,7 +70,6 @@ impl VideoCapture { .await; }); - // In the _queue_thread let _queue_thread = tokio::spawn(async move { // Helper function to push to queue and handle errors @@ -129,7 +129,7 @@ impl VideoCapture { &output_path, fps, new_chunk_callback_clone, - monitor_id, + capture_source, video_chunk_duration, ) .await; @@ -170,7 +170,7 @@ pub async fn start_ffmpeg_process(output_file: &str, fps: f64) -> Result, - monitor_id: u32, + capture_source: CaptureSource<'static>, video_chunk_duration: Duration, ) { debug!("Starting save_frames_as_video function"); @@ -230,7 +230,7 @@ async fn save_frames_as_video( let first_frame = wait_for_first_frame(frame_queue).await; let buffer = encode_frame(&first_frame); - let output_file = create_output_file(output_path, monitor_id); + let output_file = create_output_file(output_path, capture_source); new_chunk_callback(&output_file); match start_ffmpeg_process(&output_file, fps).await { @@ -289,11 +289,11 @@ fn encode_frame(frame: &CaptureResult) -> Vec { buffer } -fn create_output_file(output_path: &str, monitor_id: u32) -> String { +fn create_output_file(output_path: &str, capture_source: CaptureSource) -> String { let time = Utc::now(); let formatted_time = time.format("%Y-%m-%d_%H-%M-%S").to_string(); PathBuf::from(output_path) - .join(format!("monitor_{}_{}.mp4", monitor_id, formatted_time)) + .join(format!("{}_{}.mp4", capture_source, formatted_time)) .to_str() .expect("Failed to create valid path") .to_string() diff --git a/screenpipe-vision/Cargo.toml b/screenpipe-vision/Cargo.toml index 65197c0401..f4d0c7c5c6 100644 --- a/screenpipe-vision/Cargo.toml +++ b/screenpipe-vision/Cargo.toml @@ -43,6 +43,8 @@ tracing = { workspace = true } once_cell = "1.19" which = "6.0" +futures = "0.3" + [dev-dependencies] tempfile = "3.3.0" criterion = { workspace = true } diff --git a/screenpipe-vision/src/bin/screenpipe-vision.rs b/screenpipe-vision/src/bin/screenpipe-vision.rs index 22198b456b..c7d4b1e86d 100644 --- a/screenpipe-vision/src/bin/screenpipe-vision.rs +++ b/screenpipe-vision/src/bin/screenpipe-vision.rs @@ -1,6 +1,8 @@ use clap::Parser; use screenpipe_core::Language; -use screenpipe_vision::{continuous_capture, monitor::get_default_monitor, OcrEngine}; +use screenpipe_vision::{ + continuous_capture, core::CaptureSource, monitor::get_default_monitor, OcrEngine, +}; use std::time::Duration; use tokio::sync::mpsc::channel; use tracing_subscriber::{fmt::format::FmtSpan, EnvFilter}; @@ -18,6 +20,12 @@ struct Cli { #[arg(short = 'l', long, value_enum)] language: Vec, + + #[arg(long)] + rdp_session_id: Option, + + #[arg(long)] + monitor_id: Option, } #[tokio::main] @@ -37,19 +45,24 @@ async fn main() { let save_text_files = cli.save_text_files; let languages = cli.language; - let monitor = get_default_monitor().await; - let id = monitor.id(); - tokio::spawn(async move { + let capture_source = if let Some(session_id) = cli.rdp_session_id { + CaptureSource::RdpSession(Box::leak(session_id.into_boxed_str())) + } else if let Some(monitor_id) = cli.monitor_id { + CaptureSource::LocalMonitor(monitor_id) + } else { + CaptureSource::LocalMonitor(get_default_monitor().await.id()) + }; + continuous_capture( result_tx, Duration::from_secs_f32(1.0 / cli.fps), save_text_files, OcrEngine::AppleNative, - id, + capture_source, &[], &[], - languages.clone(), + languages, ) .await }); diff --git a/screenpipe-vision/src/core.rs b/screenpipe-vision/src/core.rs index 77be55d5d9..212b68633b 100644 --- a/screenpipe-vision/src/core.rs +++ b/screenpipe-vision/src/core.rs @@ -14,13 +14,31 @@ use log::{debug, error}; use screenpipe_core::Language; use screenpipe_integrations::unstructured_ocr::perform_ocr_cloud; use serde_json; +use std::future::Future; +use std::pin::Pin; use std::{ collections::HashMap, + fmt, time::{Duration, Instant}, }; use tokio::sync::mpsc::Sender; use xcap::Monitor; +#[derive(Debug, Clone, Copy)] +pub enum CaptureSource<'a> { + LocalMonitor(u32), + RdpSession(&'a str), +} + +impl<'a> fmt::Display for CaptureSource<'a> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + CaptureSource::LocalMonitor(id) => write!(f, "monitor_{}", id), + CaptureSource::RdpSession(session) => write!(f, "rdp_{}", session), + } + } +} + pub struct CaptureResult { pub image: DynamicImage, pub frame_number: u64, @@ -51,47 +69,78 @@ pub async fn continuous_capture( interval: Duration, save_text_files_flag: bool, ocr_engine: OcrEngine, - monitor_id: u32, + capture_source: CaptureSource<'static>, ignore_list: &[String], include_list: &[String], languages: Vec, ) { debug!( - "continuous_capture: Starting using monitor: {:?}", - monitor_id + "continuous_capture: Starting with source: {:?}", + capture_source ); + + let capture_fn: Box< + dyn Fn() -> Pin< + Box< + dyn Future< + Output = Option<( + DynamicImage, + Vec<(DynamicImage, String, String, bool)>, + u64, + Duration, + )>, + > + Send, + >, + > + Send + + Sync, + > = match capture_source { + CaptureSource::LocalMonitor(monitor_id) => { + let monitor = match get_monitor_by_id(monitor_id).await { + Some(m) => m, + None => { + error!("Failed to get monitor with id: {}", monitor_id); + return; + } + }; + let ignore_list = ignore_list.to_vec(); + let include_list = include_list.to_vec(); + + Box::new(move || { + let monitor_clone = monitor.clone(); + let ignore_list = ignore_list.clone(); + let include_list = include_list.clone(); + + Box::pin(async move { + match capture_screenshot(&monitor_clone, &ignore_list, &include_list).await { + Ok(result) => Some(result), + Err(_) => None, + } + }) + }) + } + CaptureSource::RdpSession(session_id) => Box::new(move || { + Box::pin(async move { + match capture_rdp_session(&session_id).await { + Ok(image) => { + let window_images = vec![]; + let image_hash = 0; + Some((image, window_images, image_hash, Duration::from_secs(0))) + } + Err(_) => None, + } + }) + }), + }; + let mut frame_counter: u64 = 0; let mut previous_image: Option = None; let mut max_average: Option = None; let mut max_avg_value = 0.0; - let monitor = match get_monitor_by_id(monitor_id).await { - Some(m) => m, - None => { - error!( - "Failed to get monitor with id: {}. Exiting continuous_capture.", - monitor_id - ); - return; - } - }; - loop { - let capture_result = match capture_screenshot(&monitor, &ignore_list, &include_list).await { - Ok((image, window_images, image_hash, _capture_duration)) => { - debug!( - "Captured screenshot on monitor {} with hash: {}", - monitor_id, image_hash - ); - Some((image, window_images, image_hash)) - } - Err(e) => { - error!("Failed to capture screenshot: {}", e); - None - } - }; + let capture_result = capture_fn().await; - if let Some((image, window_images, image_hash)) = capture_result { + if let Some((image, window_images, image_hash, _capture_duration)) = capture_result { let current_average = match compare_with_previous_image( previous_image.as_ref(), &image, @@ -305,3 +354,33 @@ pub fn trigger_screen_capture_permission() -> Result<()> { Ok(()) } +#[allow(unused)] +pub async fn capture_rdp_session(session_id: &str) -> Result { + #[cfg(target_os = "windows")] + { + // windows terminal services api + use windows::Win32::System::TerminalServices::{ + WTSQuerySessionInformation, WTSVirtualChannelQuery, WTS_CURRENT_SERVER_HANDLE, + }; + + // capture the rdp session buffer + // this is a simplified version - you'll need proper error handling + let buffer = unsafe { + WTSVirtualChannelQuery( + WTS_CURRENT_SERVER_HANDLE, + session_id, + WTSVirtualChannelQuery::WTSVirtualChannelGetData, + std::ptr::null_mut(), + )? + }; + + // convert buffer to image + // you'll need to implement the actual conversion based on your needs + DynamicImage::from_buffer(&buffer) + } + + #[cfg(not(target_os = "windows"))] + { + Err(anyhow!("rdp capture only supported on windows")) + } +} diff --git a/screenpipe-vision/src/lib.rs b/screenpipe-vision/src/lib.rs index d277160f1f..f5be58bd32 100644 --- a/screenpipe-vision/src/lib.rs +++ b/screenpipe-vision/src/lib.rs @@ -1,10 +1,10 @@ #[cfg(target_os = "macos")] pub mod apple; pub mod core; -pub mod run_ui_monitoring_macos; #[cfg(target_os = "windows")] pub mod microsoft; pub mod monitor; +pub mod run_ui_monitoring_macos; pub mod tesseract; pub mod utils; #[cfg(target_os = "macos")] @@ -12,7 +12,7 @@ pub use apple::{parse_apple_ocr_result, perform_ocr_apple}; pub use core::{continuous_capture, process_ocr_task, CaptureResult}; pub use utils::OcrEngine; pub mod capture_screenshot_by_window; -pub use run_ui_monitoring_macos::run_ui; #[cfg(target_os = "windows")] pub use microsoft::perform_ocr_windows; +pub use run_ui_monitoring_macos::run_ui; pub use tesseract::perform_ocr_tesseract; From 7d995947403b077bb5ef8cfb46a9d7396b3d5051 Mon Sep 17 00:00:00 2001 From: Louis Beaumont Date: Mon, 18 Nov 2024 10:46:17 -0800 Subject: [PATCH 2/8] add example --- screenpipe-vision/examples/rdp.rs | 89 +++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 screenpipe-vision/examples/rdp.rs diff --git a/screenpipe-vision/examples/rdp.rs b/screenpipe-vision/examples/rdp.rs new file mode 100644 index 0000000000..b4cd226f85 --- /dev/null +++ b/screenpipe-vision/examples/rdp.rs @@ -0,0 +1,89 @@ +#[cfg(target_os = "windows")] +async fn capture_session(session_id: u32) -> anyhow::Result<()> { + use anyhow::Result; + use std::ptr::null_mut; + use tokio::sync::mpsc; + use windows::Win32::System::RemoteDesktop::{ + WTSActive, WTSEnumerateSessionsW, WTSVirtualChannelClose, WTSVirtualChannelOpen, + WTSVirtualChannelRead, WTSVirtualChannelWrite, WTS_CURRENT_SERVER_HANDLE, + WTS_SESSION_INFOW, + }; + + const CHANNEL_NAME: &str = "SCREENCAP\0"; + const BUFFER_SIZE: u32 = 65536; + + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = null_mut(); + + // Get all sessions + let success = unsafe { + WTSEnumerateSessionsW( + WTS_CURRENT_SERVER_HANDLE, + 0, + 1, + &mut sessions, + &mut session_count, + ) + }; + + if success == 0 { + return Err(anyhow::anyhow!("failed to enumerate sessions")); + } + + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; + + for session in sessions_slice { + if session.State == WTSActive { + let channel = unsafe { + WTSVirtualChannelOpen( + WTS_CURRENT_SERVER_HANDLE, + session.SessionId, + CHANNEL_NAME.as_ptr() as *mut i8, + ) + }; + + if channel.is_null() { + println!("failed to open channel for session {}", session.SessionId); + continue; + } + + let mut buffer = vec![0u8; BUFFER_SIZE as usize]; + let mut bytes_read: u32 = 0; + + loop { + let success = unsafe { + WTSVirtualChannelRead( + channel, + 0, + buffer.as_mut_ptr() as *mut i8, + BUFFER_SIZE, + &mut bytes_read, + ) + }; + + if success == 0 { + break; + } + + if bytes_read > 0 { + println!( + "session {}: received {} bytes", + session.SessionId, bytes_read + ); + } + } + + unsafe { WTSVirtualChannelClose(channel) }; + } + } + + Ok(()) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + #[cfg(target_os = "windows")] + capture_session(1).await; + + Ok(()) +} From 58c17407a816da2e884ff96dc188cd5107387268 Mon Sep 17 00:00:00 2001 From: Louis Beaumont Date: Mon, 18 Nov 2024 17:24:41 -0800 Subject: [PATCH 3/8] RDP works! --- screenpipe-vision/Cargo.toml | 7 + screenpipe-vision/examples/rdp.rs | 322 +++++++++++++++++++++++++----- 2 files changed, 278 insertions(+), 51 deletions(-) diff --git a/screenpipe-vision/Cargo.toml b/screenpipe-vision/Cargo.toml index f4d0c7c5c6..f9c7285005 100644 --- a/screenpipe-vision/Cargo.toml +++ b/screenpipe-vision/Cargo.toml @@ -86,6 +86,13 @@ harness = false name = "screenpipe-vision-websocket" path = "examples/websocket.rs" +[[example]] +name = "rdp" +path = "examples/rdp.rs" +required-features = ["rdp-features"] + +[features] +rdp-features = ["windows/Win32_Graphics_Gdi", "windows/Win32_Graphics"] [target.'cfg(target_os = "windows")'.dependencies] windows = { version = "0.58", features = ["Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams"] } diff --git a/screenpipe-vision/examples/rdp.rs b/screenpipe-vision/examples/rdp.rs index b4cd226f85..baa421010b 100644 --- a/screenpipe-vision/examples/rdp.rs +++ b/screenpipe-vision/examples/rdp.rs @@ -1,89 +1,309 @@ #[cfg(target_os = "windows")] -async fn capture_session(session_id: u32) -> anyhow::Result<()> { - use anyhow::Result; - use std::ptr::null_mut; - use tokio::sync::mpsc; - use windows::Win32::System::RemoteDesktop::{ - WTSActive, WTSEnumerateSessionsW, WTSVirtualChannelClose, WTSVirtualChannelOpen, - WTSVirtualChannelRead, WTSVirtualChannelWrite, WTS_CURRENT_SERVER_HANDLE, - WTS_SESSION_INFOW, - }; - - const CHANNEL_NAME: &str = "SCREENCAP\0"; - const BUFFER_SIZE: u32 = 65536; +use windows::Win32::System::RemoteDesktop::{ + WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, + WTS_INFO_CLASS, WTS_SESSION_INFOW, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Graphics::Gdi::{ + BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, + GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Foundation::HANDLE; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; + +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ + GetTokenInformation, LookupPrivilegeNameW, TokenPrivileges, SE_PRIVILEGE_ENABLED, + TOKEN_PRIVILEGES, TOKEN_QUERY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Memory::{LocalAlloc, LPTR}; + +#[cfg(target_os = "windows")] +use windows::core::PWSTR; + +use std::ffi::c_void; +use std::fs; +use std::path::Path; + +#[cfg(target_os = "windows")] +async fn capture_all_sessions() -> anyhow::Result<()> { + // Create screenshots directory if it doesn't exist + let screenshots_dir = Path::new("screenshots"); + if !screenshots_dir.exists() { + println!("creating screenshots directory..."); + fs::create_dir_all(screenshots_dir)?; + } let mut session_count: u32 = 0; - let mut sessions: *mut WTS_SESSION_INFOW = null_mut(); + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); - // Get all sessions - let success = unsafe { + unsafe { WTSEnumerateSessionsW( WTS_CURRENT_SERVER_HANDLE, 0, 1, &mut sessions, &mut session_count, - ) - }; - - if success == 0 { - return Err(anyhow::anyhow!("failed to enumerate sessions")); + )?; } let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; + println!("found {} sessions", session_count); + for session in sessions_slice { if session.State == WTSActive { - let channel = unsafe { - WTSVirtualChannelOpen( + println!("processing session {}", session.SessionId); + + // Get session username to verify we have access + let mut bytes_returned: u32 = 0; + let mut username = PWSTR::null(); + + unsafe { + match WTSQuerySessionInformationW( WTS_CURRENT_SERVER_HANDLE, session.SessionId, - CHANNEL_NAME.as_ptr() as *mut i8, + WTS_INFO_CLASS(5), // WTSUserName + &mut username, + &mut bytes_returned, + ) { + Ok(_) => { + let username_str = username.to_string()?; + println!( + "session {} belongs to user: {}", + session.SessionId, username_str + ); + } + Err(e) => { + println!( + "warning: couldn't get username for session {}: {:?}", + session.SessionId, e + ); + } + } + } + + // Create DC for this session + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, ) }; - if channel.is_null() { - println!("failed to open channel for session {}", session.SessionId); - continue; - } + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) + }; + + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - let mut buffer = vec![0u8; BUFFER_SIZE as usize]; - let mut bytes_read: u32 = 0; - - loop { - let success = unsafe { - WTSVirtualChannelRead( - channel, - 0, - buffer.as_mut_ptr() as *mut i8, - BUFFER_SIZE, - &mut bytes_read, - ) + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() }; - if success == 0 { - break; - } + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); - if bytes_read > 0 { - println!( - "session {}: received {} bytes", - session.SessionId, bytes_read - ); - } + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); } + } else { + println!("skipping inactive session {}", session.SessionId); + } + } - unsafe { WTSVirtualChannelClose(channel) }; + Ok(()) +} + +#[cfg(target_os = "windows")] +fn save_buffer_as_png( + buffer: &[u8], + width: u32, + height: u32, + filename: &str, +) -> anyhow::Result<()> { + use image::{ImageBuffer, Rgba}; + + let mut img = ImageBuffer::new(width, height); + + for (x, y, pixel) in img.enumerate_pixels_mut() { + let pos = ((y * width + x) * 4) as usize; + *pixel = Rgba([ + buffer[pos + 2], // B + buffer[pos + 1], // G + buffer[pos], // R + buffer[pos + 3], // A + ]); + } + + img.save(filename)?; + println!("saved capture to {}", filename); + Ok(()) +} + +#[cfg(target_os = "windows")] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("checking rdp permissions..."); + + unsafe { + let mut token_handle = HANDLE::default(); + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + &mut token_handle as *mut HANDLE, + )?; + println!("successfully opened process token"); + + let mut return_length = 0; + let result = + GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); + + if result.is_err() { + println!("got required buffer size: {} bytes", return_length); + + // Allocate the required buffer + let buffer = LocalAlloc(LPTR, return_length as usize)?; + let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; + + // Second call with properly sized buffer + GetTokenInformation( + token_handle, + TokenPrivileges, + Some(buffer.0 as *mut c_void), + return_length, + &mut return_length, + )?; + + let privileges = &*privileges_ptr; + println!("found {} privileges", privileges.PrivilegeCount); + + let privilege_array = std::slice::from_raw_parts( + privileges.Privileges.as_ptr(), + privileges.PrivilegeCount as usize, + ); + + for privilege in privilege_array { + // First call to get the required name length + let mut name_len = 0; + let name_result = + LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); + + // Ignore the expected error from getting the size + if name_len > 0 { + // Allocate buffer with the correct size (+1 for null terminator) + let mut name = vec![0u16; name_len as usize + 1]; + let mut final_len = name_len; + + // Second call to actually get the name + match LookupPrivilegeNameW( + None, + &privilege.Luid, + PWSTR(name.as_mut_ptr()), + &mut final_len, + ) { + Ok(_) => { + let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); + println!( + "privilege: {} (enabled: {})", + priv_name, + privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED + ); + } + Err(e) => { + println!("failed to get privilege name: {:?}", e); + } + } + } + } } } Ok(()) } +#[cfg(not(target_os = "windows"))] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("rdp permissions check only available on windows"); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +async fn capture_all_sessions() -> anyhow::Result<()> { + println!("rdp capture only available on windows"); + Ok(()) +} + #[tokio::main] async fn main() -> anyhow::Result<()> { - #[cfg(target_os = "windows")] - capture_session(1).await; + println!("starting rdp example..."); + + println!("checking permissions..."); + let err = check_rdp_permissions().await; + if let Err(e) = err { + println!("error checking permissions: {:?}", e); + } + + println!("starting capture..."); + let err = capture_all_sessions().await; + if let Err(e) = err { + println!("error capturing session: {:?}", e); + } + println!("example completed successfully"); Ok(()) } From 2cde8c132ecd1eb95a33fe143c27c88aae8bc6aa Mon Sep 17 00:00:00 2001 From: Louis Beaumont Date: Mon, 18 Nov 2024 17:47:38 -0800 Subject: [PATCH 4/8] chore: polish integration --- .../src/bin/screenpipe-server.rs | 6 +- screenpipe-server/src/cli.rs | 4 +- screenpipe-vision/src/core.rs | 31 +-- screenpipe-vision/src/lib.rs | 1 + screenpipe-vision/src/remote_desktop.rs | 218 ++++++++++++++++++ 5 files changed, 225 insertions(+), 35 deletions(-) create mode 100644 screenpipe-vision/src/remote_desktop.rs diff --git a/screenpipe-server/src/bin/screenpipe-server.rs b/screenpipe-server/src/bin/screenpipe-server.rs index fe5a940c9c..bc3228ada1 100644 --- a/screenpipe-server/src/bin/screenpipe-server.rs +++ b/screenpipe-server/src/bin/screenpipe-server.rs @@ -15,7 +15,7 @@ use screenpipe_core::find_ffmpeg_path; use screenpipe_server::{ cli::{Cli, CliAudioTranscriptionEngine, CliOcrEngine, Command, PipeCommand}, start_continuous_recording, watch_pid, DatabaseManager, PipeControl, PipeManager, ResourceMonitor, Server, highlight::{Highlight,HighlightConfig} }; -use screenpipe_vision::{core::CaptureSource, monitor::{get_default_monitor, list_monitors}}; +use screenpipe_vision::{core::CaptureSource, monitor::{get_default_monitor, list_monitors}, remote_desktop::list_rdp_sessions}; #[cfg(target_os = "macos")] use screenpipe_vision::run_ui; use serde_json::{json, Value}; @@ -365,8 +365,8 @@ async fn main() -> anyhow::Result<()> { loop { let vad_engine_clone = vad_engine.clone(); // Clone it here for each iteration let mut shutdown_rx = shutdown_tx_clone.subscribe(); - let capture_source = if cli.rdp_session_id.len() > 0 { - cli.rdp_session_id.iter().map(|s| CaptureSource::RdpSession(Box::leak(s.clone().into_boxed_str()))).collect() + let capture_source = if cli.use_remote_desktop { + list_rdp_sessions().await.unwrap().iter().map(|(id, _username)| CaptureSource::RdpSession(Box::leak(id.to_string().into_boxed_str()))).collect() } else if cli.monitor_id.len() > 0 { cli.monitor_id.iter().map(|id| CaptureSource::LocalMonitor(*id)).collect() } else { diff --git a/screenpipe-server/src/cli.rs b/screenpipe-server/src/cli.rs index f1b97a83d7..49fcc6fc9a 100644 --- a/screenpipe-server/src/cli.rs +++ b/screenpipe-server/src/cli.rs @@ -243,9 +243,9 @@ pub struct Cli { #[arg(long, default_value_t = false)] pub enable_frame_cache: bool, - /// RDP session ID to use for screen capture - this will override the monitor ID + /// If enabled will capture all RDP sessions on the machine #[arg(long)] - pub rdp_session_id: Vec, + pub use_remote_desktop: bool, #[command(subcommand)] pub command: Option, diff --git a/screenpipe-vision/src/core.rs b/screenpipe-vision/src/core.rs index 212b68633b..eb678eceae 100644 --- a/screenpipe-vision/src/core.rs +++ b/screenpipe-vision/src/core.rs @@ -5,6 +5,7 @@ use crate::apple::perform_ocr_apple; #[cfg(target_os = "windows")] use crate::microsoft::perform_ocr_windows; use crate::monitor::get_monitor_by_id; +use crate::remote_desktop::capture_rdp_session; use crate::tesseract::perform_ocr_tesseract; use crate::utils::OcrEngine; use crate::utils::{capture_screenshot, compare_with_previous_image, save_text_files}; @@ -354,33 +355,3 @@ pub fn trigger_screen_capture_permission() -> Result<()> { Ok(()) } -#[allow(unused)] -pub async fn capture_rdp_session(session_id: &str) -> Result { - #[cfg(target_os = "windows")] - { - // windows terminal services api - use windows::Win32::System::TerminalServices::{ - WTSQuerySessionInformation, WTSVirtualChannelQuery, WTS_CURRENT_SERVER_HANDLE, - }; - - // capture the rdp session buffer - // this is a simplified version - you'll need proper error handling - let buffer = unsafe { - WTSVirtualChannelQuery( - WTS_CURRENT_SERVER_HANDLE, - session_id, - WTSVirtualChannelQuery::WTSVirtualChannelGetData, - std::ptr::null_mut(), - )? - }; - - // convert buffer to image - // you'll need to implement the actual conversion based on your needs - DynamicImage::from_buffer(&buffer) - } - - #[cfg(not(target_os = "windows"))] - { - Err(anyhow!("rdp capture only supported on windows")) - } -} diff --git a/screenpipe-vision/src/lib.rs b/screenpipe-vision/src/lib.rs index f5be58bd32..81ef0fabb4 100644 --- a/screenpipe-vision/src/lib.rs +++ b/screenpipe-vision/src/lib.rs @@ -4,6 +4,7 @@ pub mod core; #[cfg(target_os = "windows")] pub mod microsoft; pub mod monitor; +pub mod remote_desktop; pub mod run_ui_monitoring_macos; pub mod tesseract; pub mod utils; diff --git a/screenpipe-vision/src/remote_desktop.rs b/screenpipe-vision/src/remote_desktop.rs new file mode 100644 index 0000000000..62d45a3061 --- /dev/null +++ b/screenpipe-vision/src/remote_desktop.rs @@ -0,0 +1,218 @@ +use image::DynamicImage; + +// TODO: this file could work for things like macos client connecting on a linux or stuff like that (tmux?) + +// TODO: atm ignored windows and included windows not supported in RDP, not sure about multiple monitors + +#[allow(unused)] +pub async fn capture_rdp_session(session_id: &str) -> anyhow::Result { + #[cfg(target_os = "windows")] + { + use image::{ImageBuffer, Rgba}; + use windows::Win32::Graphics::Gdi::{ + BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, + GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, + }; + + // Create DC for this session + let dc_name = format!("DISPLAY#{}\0", session_id); + println!("creating dc with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to make this configurable + 1080, // height - you might want to make this configurable + ) + }; + + let mut buffer = unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + + buffer + }; + + // Convert buffer to DynamicImage + let mut img = ImageBuffer::new(1920, 1080); + for (x, y, pixel) in img.enumerate_pixels_mut() { + let pos = ((y * 1920 + x) * 4) as usize; + *pixel = Rgba([ + buffer[pos + 2], // B + buffer[pos + 1], // G + buffer[pos], // R + buffer[pos + 3], // A + ]); + } + + return Ok(DynamicImage::ImageRgba8(img)); + } + + #[cfg(not(target_os = "windows"))] + { + anyhow::bail!("rdp capture only supported on windows") + } +} + +#[cfg(target_os = "windows")] +pub async fn list_rdp_sessions() -> anyhow::Result> { + use windows::core::PWSTR; + use windows::Win32::System::RemoteDesktop::{ + WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, + WTS_INFO_CLASS, WTS_SESSION_INFOW, + }; + + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + let mut result = Vec::new(); + + unsafe { + WTSEnumerateSessionsW( + WTS_CURRENT_SERVER_HANDLE, + 0, + 1, + &mut sessions, + &mut session_count, + )?; + + let sessions_slice = std::slice::from_raw_parts(sessions, session_count as usize); + + for session in sessions_slice { + if session.State == WTSActive { + let mut bytes_returned: u32 = 0; + let mut username = PWSTR::null(); + + match WTSQuerySessionInformationW( + WTS_CURRENT_SERVER_HANDLE, + session.SessionId, + WTS_INFO_CLASS(5), // WTSUserName + &mut username, + &mut bytes_returned, + ) { + Ok(_) => { + let username_str = username.to_string()?; + result.push((session.SessionId, username_str)); + } + Err(_) => { + result.push((session.SessionId, String::from("unknown"))); + } + } + } + } + } + + Ok(result) +} + +#[cfg(not(target_os = "windows"))] +pub async fn list_rdp_sessions() -> anyhow::Result> { + anyhow::bail!("rdp session listing only supported on windows") +} + +#[cfg(target_os = "windows")] +pub struct SessionInfo { + pub session_id: u32, + pub username: String, + pub domain: String, + pub client_name: String, + pub client_address: String, + pub connection_state: String, +} + +#[cfg(target_os = "windows")] +pub async fn get_session_details(session_id: u32) -> anyhow::Result { + use windows::core::PWSTR; + use windows::Win32::System::RemoteDesktop::{ + WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, WTS_INFO_CLASS, + }; + + unsafe { + let mut query_info = |info_class: WTS_INFO_CLASS| -> anyhow::Result { + let mut bytes_returned: u32 = 0; + let mut buffer = PWSTR::null(); + + WTSQuerySessionInformationW( + WTS_CURRENT_SERVER_HANDLE, + session_id, + info_class, + &mut buffer, + &mut bytes_returned, + )?; + + Ok(buffer + .to_string() + .unwrap_or_else(|_| String::from("unknown"))) + }; + + let info = SessionInfo { + session_id, + username: query_info(WTS_INFO_CLASS(5))?, // WTSUserName + domain: query_info(WTS_INFO_CLASS(7))?, // WTSDomainName + client_name: query_info(WTS_INFO_CLASS(10))?, // WTSClientName + client_address: query_info(WTS_INFO_CLASS(14))?, // WTSClientAddress + connection_state: query_info(WTS_INFO_CLASS(8))?, // WTSConnectState + }; + + Ok(info) + } +} + +#[cfg(not(target_os = "windows"))] +pub struct SessionInfo { + pub session_id: u32, + pub username: String, + pub domain: String, + pub client_name: String, + pub client_address: String, + pub connection_state: String, +} + +#[cfg(not(target_os = "windows"))] +pub async fn get_session_details(_session_id: u32) -> anyhow::Result { + anyhow::bail!("session details only supported on windows") +} From 23133b94f3cd729ae4d1f072c2256079dc45c637 Mon Sep 17 00:00:00 2001 From: louis030195 Date: Thu, 21 Nov 2024 00:32:55 +0000 Subject: [PATCH 5/8] make it work --- screenpipe-vision/Cargo.toml | 6 +- screenpipe-vision/bin | 316 +++++++++++++++++++ screenpipe-vision/examples/rdp.rs | 239 +++++++------- screenpipe-vision/src/bin/rdp.rs | 316 +++++++++++++++++++ screenpipe-vision/src/core.rs | 28 +- screenpipe-vision/src/remote_desktop.rs | 398 +++++++++++++++++++++--- 6 files changed, 1124 insertions(+), 179 deletions(-) create mode 100644 screenpipe-vision/bin create mode 100644 screenpipe-vision/src/bin/rdp.rs diff --git a/screenpipe-vision/Cargo.toml b/screenpipe-vision/Cargo.toml index f9c7285005..58b25c29e4 100644 --- a/screenpipe-vision/Cargo.toml +++ b/screenpipe-vision/Cargo.toml @@ -89,13 +89,9 @@ path = "examples/websocket.rs" [[example]] name = "rdp" path = "examples/rdp.rs" -required-features = ["rdp-features"] - -[features] -rdp-features = ["windows/Win32_Graphics_Gdi", "windows/Win32_Graphics"] [target.'cfg(target_os = "windows")'.dependencies] -windows = { version = "0.58", features = ["Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams"] } +windows = { version = "0.58", features = ["Win32_System_ProcessStatus", "Win32_UI_WindowsAndMessaging", "Win32_UI", "Win32_Graphics_Gdi", "Win32_Graphics", "Win32_System_Memory", "Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams", "Win32_System_RemoteDesktop", "Win32_System_Diagnostics", "Win32_System_Diagnostics_Debug", "Win32_Security", "Win32_System_Threading", "Win32_System_Registry", "Win32_System_StationsAndDesktops"] } xcap = "0.0.12" [target.'cfg(target_os = "macos")'.dependencies] diff --git a/screenpipe-vision/bin b/screenpipe-vision/bin new file mode 100644 index 0000000000..4083b7a70e --- /dev/null +++ b/screenpipe-vision/bin @@ -0,0 +1,316 @@ +#[cfg(target_os = "windows")] +use windows::Win32::System::RemoteDesktop::{ + WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, + WTS_INFO_CLASS, WTS_SESSION_INFOW, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Graphics::Gdi::{ + BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, + GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Foundation::HANDLE; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; + +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ + GetTokenInformation, LookupPrivilegeNameW, TokenPrivileges, SE_PRIVILEGE_ENABLED, + TOKEN_PRIVILEGES, TOKEN_QUERY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Memory::{LocalAlloc, LPTR}; + +#[cfg(target_os = "windows")] +use windows::core::PWSTR; + +use std::ffi::c_void; +use std::fs; +use std::path::Path; + +#[cfg(target_os = "windows")] +async fn capture_all_sessions() -> anyhow::Result<()> { + // Create screenshots directory if it doesn't exist + let screenshots_dir = Path::new("screenshots"); + if !screenshots_dir.exists() { + println!("creating screenshots directory..."); + fs::create_dir_all(screenshots_dir)?; + } + + loop { + println!("capturing new round of screenshots..."); + + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + + unsafe { + WTSEnumerateSessionsW( + WTS_CURRENT_SERVER_HANDLE, + 0, + 1, + &mut sessions, + &mut session_count, + )?; + } + + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; + + println!("found {} sessions", session_count); + + for session in sessions_slice { + if session.State == WTSActive { + println!("processing session {}", session.SessionId); + + // Get session username to verify we have access + let mut bytes_returned: u32 = 0; + let mut username = PWSTR::null(); + + unsafe { + match WTSQuerySessionInformationW( + WTS_CURRENT_SERVER_HANDLE, + session.SessionId, + WTS_INFO_CLASS(5), // WTSUserName + &mut username, + &mut bytes_returned, + ) { + Ok(_) => { + let username_str = username.to_string()?; + println!( + "session {} belongs to user: {}", + session.SessionId, username_str + ); + } + Err(e) => { + println!( + "warning: couldn't get username for session {}: {:?}", + session.SessionId, e + ); + } + } + } + + // Create DC for this session + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) + }; + + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + } + } else { + println!("skipping inactive session {}", session.SessionId); + } + } + + // Add delay between captures + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + } + + Ok(()) +} + +#[cfg(target_os = "windows")] +fn save_buffer_as_png( + buffer: &[u8], + width: u32, + height: u32, + filename: &str, +) -> anyhow::Result<()> { + use image::{ImageBuffer, Rgba}; + + let mut img = ImageBuffer::new(width, height); + + for (x, y, pixel) in img.enumerate_pixels_mut() { + let pos = ((y * width + x) * 4) as usize; + *pixel = Rgba([ + buffer[pos + 2], // B + buffer[pos + 1], // G + buffer[pos], // R + buffer[pos + 3], // A + ]); + } + + img.save(filename)?; + println!("saved capture to {}", filename); + Ok(()) +} + +#[cfg(target_os = "windows")] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("checking rdp permissions..."); + + unsafe { + let mut token_handle = HANDLE::default(); + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + &mut token_handle as *mut HANDLE, + )?; + println!("successfully opened process token"); + + let mut return_length = 0; + let result = + GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); + + if result.is_err() { + println!("got required buffer size: {} bytes", return_length); + + // Allocate the required buffer + let buffer = LocalAlloc(LPTR, return_length as usize)?; + let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; + + // Second call with properly sized buffer + GetTokenInformation( + token_handle, + TokenPrivileges, + Some(buffer.0 as *mut c_void), + return_length, + &mut return_length, + )?; + + let privileges = &*privileges_ptr; + println!("found {} privileges", privileges.PrivilegeCount); + + let privilege_array = std::slice::from_raw_parts( + privileges.Privileges.as_ptr(), + privileges.PrivilegeCount as usize, + ); + + for privilege in privilege_array { + // First call to get the required name length + let mut name_len = 0; + let name_result = + LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); + + // Ignore the expected error from getting the size + if name_len > 0 { + // Allocate buffer with the correct size (+1 for null terminator) + let mut name = vec![0u16; name_len as usize + 1]; + let mut final_len = name_len; + + // Second call to actually get the name + match LookupPrivilegeNameW( + None, + &privilege.Luid, + PWSTR(name.as_mut_ptr()), + &mut final_len, + ) { + Ok(_) => { + let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); + println!( + "privilege: {} (enabled: {})", + priv_name, + privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED + ); + } + Err(e) => { + println!("failed to get privilege name: {:?}", e); + } + } + } + } + } + } + + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("rdp permissions check only available on windows"); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +async fn capture_all_sessions() -> anyhow::Result<()> { + println!("rdp capture only available on windows"); + Ok(()) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + println!("starting rdp example..."); + + println!("checking permissions..."); + let err = check_rdp_permissions().await; + if let Err(e) = err { + println!("error checking permissions: {:?}", e); + } + + println!("starting continuous capture..."); + let err = capture_all_sessions().await; + if let Err(e) = err { + println!("error capturing session: {:?}", e); + } + + println!("example completed successfully"); + Ok(()) +} diff --git a/screenpipe-vision/examples/rdp.rs b/screenpipe-vision/examples/rdp.rs index baa421010b..4083b7a70e 100644 --- a/screenpipe-vision/examples/rdp.rs +++ b/screenpipe-vision/examples/rdp.rs @@ -41,130 +41,137 @@ async fn capture_all_sessions() -> anyhow::Result<()> { fs::create_dir_all(screenshots_dir)?; } - let mut session_count: u32 = 0; - let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + loop { + println!("capturing new round of screenshots..."); + + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + + unsafe { + WTSEnumerateSessionsW( + WTS_CURRENT_SERVER_HANDLE, + 0, + 1, + &mut sessions, + &mut session_count, + )?; + } - unsafe { - WTSEnumerateSessionsW( - WTS_CURRENT_SERVER_HANDLE, - 0, - 1, - &mut sessions, - &mut session_count, - )?; - } + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - - println!("found {} sessions", session_count); - - for session in sessions_slice { - if session.State == WTSActive { - println!("processing session {}", session.SessionId); - - // Get session username to verify we have access - let mut bytes_returned: u32 = 0; - let mut username = PWSTR::null(); - - unsafe { - match WTSQuerySessionInformationW( - WTS_CURRENT_SERVER_HANDLE, - session.SessionId, - WTS_INFO_CLASS(5), // WTSUserName - &mut username, - &mut bytes_returned, - ) { - Ok(_) => { - let username_str = username.to_string()?; - println!( - "session {} belongs to user: {}", - session.SessionId, username_str - ); - } - Err(e) => { - println!( - "warning: couldn't get username for session {}: {:?}", - session.SessionId, e - ); + println!("found {} sessions", session_count); + + for session in sessions_slice { + if session.State == WTSActive { + println!("processing session {}", session.SessionId); + + // Get session username to verify we have access + let mut bytes_returned: u32 = 0; + let mut username = PWSTR::null(); + + unsafe { + match WTSQuerySessionInformationW( + WTS_CURRENT_SERVER_HANDLE, + session.SessionId, + WTS_INFO_CLASS(5), // WTSUserName + &mut username, + &mut bytes_returned, + ) { + Ok(_) => { + let username_str = username.to_string()?; + println!( + "session {} belongs to user: {}", + session.SessionId, username_str + ); + } + Err(e) => { + println!( + "warning: couldn't get username for session {}: {:?}", + session.SessionId, e + ); + } } } - } - // Create DC for this session - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() + // Create DC for this session + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) }; - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); - - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; - - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + } + } else { + println!("skipping inactive session {}", session.SessionId); } - } else { - println!("skipping inactive session {}", session.SessionId); } + + // Add delay between captures + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; } Ok(()) @@ -298,7 +305,7 @@ async fn main() -> anyhow::Result<()> { println!("error checking permissions: {:?}", e); } - println!("starting capture..."); + println!("starting continuous capture..."); let err = capture_all_sessions().await; if let Err(e) = err { println!("error capturing session: {:?}", e); diff --git a/screenpipe-vision/src/bin/rdp.rs b/screenpipe-vision/src/bin/rdp.rs new file mode 100644 index 0000000000..4083b7a70e --- /dev/null +++ b/screenpipe-vision/src/bin/rdp.rs @@ -0,0 +1,316 @@ +#[cfg(target_os = "windows")] +use windows::Win32::System::RemoteDesktop::{ + WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, + WTS_INFO_CLASS, WTS_SESSION_INFOW, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Graphics::Gdi::{ + BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, + GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::Foundation::HANDLE; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; + +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ + GetTokenInformation, LookupPrivilegeNameW, TokenPrivileges, SE_PRIVILEGE_ENABLED, + TOKEN_PRIVILEGES, TOKEN_QUERY, +}; + +#[cfg(target_os = "windows")] +use windows::Win32::System::Memory::{LocalAlloc, LPTR}; + +#[cfg(target_os = "windows")] +use windows::core::PWSTR; + +use std::ffi::c_void; +use std::fs; +use std::path::Path; + +#[cfg(target_os = "windows")] +async fn capture_all_sessions() -> anyhow::Result<()> { + // Create screenshots directory if it doesn't exist + let screenshots_dir = Path::new("screenshots"); + if !screenshots_dir.exists() { + println!("creating screenshots directory..."); + fs::create_dir_all(screenshots_dir)?; + } + + loop { + println!("capturing new round of screenshots..."); + + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + + unsafe { + WTSEnumerateSessionsW( + WTS_CURRENT_SERVER_HANDLE, + 0, + 1, + &mut sessions, + &mut session_count, + )?; + } + + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; + + println!("found {} sessions", session_count); + + for session in sessions_slice { + if session.State == WTSActive { + println!("processing session {}", session.SessionId); + + // Get session username to verify we have access + let mut bytes_returned: u32 = 0; + let mut username = PWSTR::null(); + + unsafe { + match WTSQuerySessionInformationW( + WTS_CURRENT_SERVER_HANDLE, + session.SessionId, + WTS_INFO_CLASS(5), // WTSUserName + &mut username, + &mut bytes_returned, + ) { + Ok(_) => { + let username_str = username.to_string()?; + println!( + "session {} belongs to user: {}", + session.SessionId, username_str + ); + } + Err(e) => { + println!( + "warning: couldn't get username for session {}: {:?}", + session.SessionId, e + ); + } + } + } + + // Create DC for this session + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) + }; + + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + } + } else { + println!("skipping inactive session {}", session.SessionId); + } + } + + // Add delay between captures + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + } + + Ok(()) +} + +#[cfg(target_os = "windows")] +fn save_buffer_as_png( + buffer: &[u8], + width: u32, + height: u32, + filename: &str, +) -> anyhow::Result<()> { + use image::{ImageBuffer, Rgba}; + + let mut img = ImageBuffer::new(width, height); + + for (x, y, pixel) in img.enumerate_pixels_mut() { + let pos = ((y * width + x) * 4) as usize; + *pixel = Rgba([ + buffer[pos + 2], // B + buffer[pos + 1], // G + buffer[pos], // R + buffer[pos + 3], // A + ]); + } + + img.save(filename)?; + println!("saved capture to {}", filename); + Ok(()) +} + +#[cfg(target_os = "windows")] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("checking rdp permissions..."); + + unsafe { + let mut token_handle = HANDLE::default(); + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + &mut token_handle as *mut HANDLE, + )?; + println!("successfully opened process token"); + + let mut return_length = 0; + let result = + GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); + + if result.is_err() { + println!("got required buffer size: {} bytes", return_length); + + // Allocate the required buffer + let buffer = LocalAlloc(LPTR, return_length as usize)?; + let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; + + // Second call with properly sized buffer + GetTokenInformation( + token_handle, + TokenPrivileges, + Some(buffer.0 as *mut c_void), + return_length, + &mut return_length, + )?; + + let privileges = &*privileges_ptr; + println!("found {} privileges", privileges.PrivilegeCount); + + let privilege_array = std::slice::from_raw_parts( + privileges.Privileges.as_ptr(), + privileges.PrivilegeCount as usize, + ); + + for privilege in privilege_array { + // First call to get the required name length + let mut name_len = 0; + let name_result = + LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); + + // Ignore the expected error from getting the size + if name_len > 0 { + // Allocate buffer with the correct size (+1 for null terminator) + let mut name = vec![0u16; name_len as usize + 1]; + let mut final_len = name_len; + + // Second call to actually get the name + match LookupPrivilegeNameW( + None, + &privilege.Luid, + PWSTR(name.as_mut_ptr()), + &mut final_len, + ) { + Ok(_) => { + let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); + println!( + "privilege: {} (enabled: {})", + priv_name, + privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED + ); + } + Err(e) => { + println!("failed to get privilege name: {:?}", e); + } + } + } + } + } + } + + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +async fn check_rdp_permissions() -> anyhow::Result<()> { + println!("rdp permissions check only available on windows"); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +async fn capture_all_sessions() -> anyhow::Result<()> { + println!("rdp capture only available on windows"); + Ok(()) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + println!("starting rdp example..."); + + println!("checking permissions..."); + let err = check_rdp_permissions().await; + if let Err(e) = err { + println!("error checking permissions: {:?}", e); + } + + println!("starting continuous capture..."); + let err = capture_all_sessions().await; + if let Err(e) = err { + println!("error capturing session: {:?}", e); + } + + println!("example completed successfully"); + Ok(()) +} diff --git a/screenpipe-vision/src/core.rs b/screenpipe-vision/src/core.rs index eb678eceae..174e24f723 100644 --- a/screenpipe-vision/src/core.rs +++ b/screenpipe-vision/src/core.rs @@ -119,18 +119,24 @@ pub async fn continuous_capture( }) }) } - CaptureSource::RdpSession(session_id) => Box::new(move || { - Box::pin(async move { - match capture_rdp_session(&session_id).await { - Ok(image) => { - let window_images = vec![]; - let image_hash = 0; - Some((image, window_images, image_hash, Duration::from_secs(0))) - } - Err(_) => None, - } + CaptureSource::RdpSession(session_id) => { + let session_id = session_id.to_string(); + let ignore_list = ignore_list.to_vec(); + let include_list = include_list.to_vec(); + + Box::new(move || { + let session_id = session_id.clone(); + let ignore_list = ignore_list.clone(); + let include_list = include_list.clone(); + + Box::pin(async move { + capture_rdp_session(&session_id, &ignore_list, &include_list).await + .map(|(image, window_images, image_hash, capture_duration)| + (image, window_images, image_hash, capture_duration)) + .ok() + }) }) - }), + } }; let mut frame_counter: u64 = 0; diff --git a/screenpipe-vision/src/remote_desktop.rs b/screenpipe-vision/src/remote_desktop.rs index 62d45a3061..e5ffa8e741 100644 --- a/screenpipe-vision/src/remote_desktop.rs +++ b/screenpipe-vision/src/remote_desktop.rs @@ -1,51 +1,91 @@ use image::DynamicImage; +use log::{debug, error}; +use std::time::{Duration, Instant}; +use windows::Win32::Foundation::HWND; +use windows::Win32::Graphics::Gdi::{ + BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, GetDIBits, GetDeviceCaps, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, HORZRES, SRCCOPY, VERTRES +}; +use windows::Win32::UI::WindowsAndMessaging::{ + EnumWindows, GetWindowRect, GetWindowTextW, GetWindowThreadProcessId, IsWindowVisible, +}; +use windows::Win32::System::ProcessStatus::K32GetModuleFileNameExW; +use windows::Win32::System::Threading::{PROCESS_QUERY_INFORMATION, PROCESS_VM_READ, OpenProcess}; +use windows::Win32::System::RemoteDesktop::{ + ProcessIdToSessionId, +}; +use std::collections::HashMap; -// TODO: this file could work for things like macos client connecting on a linux or stuff like that (tmux?) - -// TODO: atm ignored windows and included windows not supported in RDP, not sure about multiple monitors - -#[allow(unused)] -pub async fn capture_rdp_session(session_id: &str) -> anyhow::Result { +pub async fn capture_rdp_session( + session_id: &str, + ignore_list: &[String], + include_list: &[String], +) -> anyhow::Result<( + DynamicImage, + Vec<(DynamicImage, String, String, bool)>, + u64, + Duration, +)> { #[cfg(target_os = "windows")] { use image::{ImageBuffer, Rgba}; - use windows::Win32::Graphics::Gdi::{ - BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, - GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, - }; + + let capture_start = Instant::now(); + + // Verify session is active + let session_info = get_session_details(session_id.parse().map_err(|e| { + debug!("invalid session id '{}': {}", session_id, e); + anyhow::anyhow!("invalid session id '{}': {}", session_id, e) + })?).await.map_err(|e| { + debug!("failed to get session details for '{}': {}", session_id, e); + anyhow::anyhow!("failed to get session details for '{}': {}", session_id, e) + })?; + + if session_info.connection_state != "Active" { + debug!("rdp session '{}' is not active (state: {})", session_id, session_info.connection_state); + anyhow::bail!("rdp session not active") + } - // Create DC for this session + // Capture full screen let dc_name = format!("DISPLAY#{}\0", session_id); - println!("creating dc with name: {}", dc_name.trim_end_matches('\0')); + debug!("creating dc with name: {}", dc_name.trim_end_matches('\0')); - let hdc = unsafe { - CreateDCA( + let (full_image, screen_width, screen_height) = unsafe { + let hdc = CreateDCA( windows::core::PCSTR(b"DISPLAY\0".as_ptr()), windows::core::PCSTR(dc_name.as_bytes().as_ptr()), None, None, - ) - }; + ); - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to make this configurable - 1080, // height - you might want to make this configurable - ) - }; + if hdc.is_invalid() { + error!("failed to create display dc for rdp session '{}' - check if you have the right permissions", session_id); + anyhow::bail!("failed to create display dc for rdp session '{}' - check if you have the right permissions", session_id); + } + + let screen_width = GetDeviceCaps(hdc, HORZRES) as usize; + let screen_height = GetDeviceCaps(hdc, VERTRES) as usize; + + let hdc_mem = CreateCompatibleDC(hdc); + let hbitmap = CreateCompatibleBitmap(hdc, screen_width as i32, screen_height as i32); - let mut buffer = unsafe { SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + BitBlt( + hdc_mem, + 0, + 0, + screen_width as i32, + screen_height as i32, + hdc, + 0, + 0, + SRCCOPY, + )?; - // Setup bitmap info let mut bi = BITMAPINFO { bmiHeader: BITMAPINFOHEADER { biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down + biWidth: screen_width as i32, + biHeight: -(screen_height as i32), biPlanes: 1, biBitCount: 32, biCompression: 0, @@ -58,47 +98,298 @@ pub async fn capture_rdp_session(session_id: &str) -> anyhow::Result { + window_images.push((window_image, window_info.app_name, window_info.window_title, false)); + } + Err(e) => { + error!( + "failed to capture window '{}' ({}) in rdp session '{}': {}", + window_info.window_title, + window_info.app_name, + session_id, + e + ); + // Continue with other windows even if one fails + continue; + } + } } - return Ok(DynamicImage::ImageRgba8(img)); + if window_images.is_empty() { + debug!("no windows were captured for rdp session '{}'", session_id); + } + + let image_hash = crate::utils::calculate_hash(&full_image); + let capture_duration = capture_start.elapsed(); + + Ok((full_image, window_images, image_hash, capture_duration)) } #[cfg(not(target_os = "windows"))] { - anyhow::bail!("rdp capture only supported on windows") + error!("rdp capture is only supported on windows"); + anyhow::bail!("rdp capture is only supported on windows") } } + + +#[cfg(target_os = "windows")] +pub async fn list_rdp_windows(target_session_id: u32) -> anyhow::Result> { + let windows_list = Vec::new(); + + unsafe extern "system" fn enum_window_callback( + hwnd: HWND, + lparam: windows::Win32::Foundation::LPARAM, + ) -> windows::Win32::Foundation::BOOL { + let windows_list = &mut *(lparam.0 as *mut Vec); + let mut process_id: u32 = 0; + GetWindowThreadProcessId(hwnd, Some(&mut process_id)); + + let mut session_id: u32 = 0; + if ProcessIdToSessionId(process_id, &mut session_id).is_ok() { + if session_id == *(lparam.0 as *const u32).add(1) && IsWindowVisible(hwnd).as_bool() { + let mut title = [0u16; 512]; + let len = GetWindowTextW(hwnd, &mut title); + if len > 0 { + let window_title = String::from_utf16_lossy(&title[..len as usize]); + + let process_name = { + let process_handle = OpenProcess( + PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, + false, + process_id, + ); + + if let Ok(handle) = process_handle { + let mut buffer = [0u16; 260]; // MAX_PATH + let len = K32GetModuleFileNameExW(handle, None, &mut buffer); + if len > 0 { + let path = String::from_utf16_lossy(&buffer[..len as usize]); + path.split('\\').last().unwrap_or("unknown").to_string() + } else { + "unknown".to_string() + } + } else { + "unknown".to_string() + } + }; + + let window_id = hwnd.0 as u64; + WINDOW_HANDLES.with(|handles| { + handles.borrow_mut().insert(window_id, hwnd.0); + }); + + windows_list.push(WindowInfo { + window_id, + app_name: process_name, + window_title, + }); + } + } + } + windows::Win32::Foundation::BOOL(1) + } + + // Create a struct to hold both the windows_list and target_session_id + #[repr(C)] + struct EnumWindowsData { + windows_list: Vec, + target_session_id: u32, + } + + let mut data = EnumWindowsData { + windows_list, + target_session_id, + }; + + unsafe { + EnumWindows( + Some(enum_window_callback), + windows::Win32::Foundation::LPARAM(&mut data as *mut _ as isize), + )?; + + Ok(data.windows_list) + } +} + +#[cfg(target_os = "windows")] +async fn capture_rdp_window( + session_id: &str, + window_id: u64, +) -> anyhow::Result { + use windows::Win32::Foundation::RECT; + + let hwnd = WINDOW_HANDLES.with(|handles| { + handles + .borrow() + .get(&window_id) + .copied() + .ok_or_else(|| anyhow::anyhow!("Window handle not found")) + })?; + + unsafe { + let hwnd = HWND(hwnd); + let mut rect: RECT = std::mem::zeroed(); + GetWindowRect(hwnd, &mut rect)?; + + let width = rect.right - rect.left; + let height = rect.bottom - rect.top; + + if width <= 0 || height <= 0 { + anyhow::bail!("Invalid window dimensions"); + } + + let dc_name = format!("DISPLAY#{}\0", session_id); + let hdc = CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ); + + let hdc_mem = CreateCompatibleDC(hdc); + let hbitmap = CreateCompatibleBitmap(hdc, width, height); + + SelectObject(hdc_mem, hbitmap); + BitBlt( + hdc_mem, + 0, + 0, + width, + height, + hdc, + rect.left, + rect.top, + SRCCOPY, + )?; + + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: width, + biHeight: -height, + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + let mut buffer = vec![0u8; (width * height * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + height as u32, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + + let mut img = image::ImageBuffer::new(width as u32, height as u32); + for (x, y, pixel) in img.enumerate_pixels_mut() { + let x = x as usize; + let y = y as usize; + let pos = (y * width as usize + x) * 4; + *pixel = image::Rgba([ + buffer[pos + 2], + buffer[pos + 1], + buffer[pos], + buffer[pos + 3], + ]); + } + + Ok(DynamicImage::ImageRgba8(img)) + } +} + +fn should_capture_window( + app_name: &str, + window_title: &str, + ignore_list: &[String], + include_list: &[String], +) -> bool { + let app_name_lower = app_name.to_lowercase(); + let title_lower = window_title.to_lowercase(); + + // Check ignore list + if ignore_list.iter().any(|ignore| { + let ignore_lower = ignore.to_lowercase(); + app_name_lower.contains(&ignore_lower) || title_lower.contains(&ignore_lower) + }) { + return false; + } + + // If include list is empty, capture all non-ignored windows + if include_list.is_empty() { + return true; + } + + // Check include list + include_list.iter().any(|include| { + let include_lower = include.to_lowercase(); + app_name_lower.contains(&include_lower) || title_lower.contains(&include_lower) + }) +} + #[cfg(target_os = "windows")] pub async fn list_rdp_sessions() -> anyhow::Result> { use windows::core::PWSTR; @@ -172,7 +463,7 @@ pub async fn get_session_details(session_id: u32) -> anyhow::Result }; unsafe { - let mut query_info = |info_class: WTS_INFO_CLASS| -> anyhow::Result { + let query_info = |info_class: WTS_INFO_CLASS| -> anyhow::Result { let mut bytes_returned: u32 = 0; let mut buffer = PWSTR::null(); @@ -216,3 +507,16 @@ pub struct SessionInfo { pub async fn get_session_details(_session_id: u32) -> anyhow::Result { anyhow::bail!("session details only supported on windows") } + +// Update the WindowInfo struct to store the raw pointer value +#[derive(Clone)] +pub struct WindowInfo { + pub window_id: u64, // Use a numeric ID instead of raw pointer + pub app_name: String, + pub window_title: String, +} + +// Keep track of window handles in a thread-local map +thread_local! { + static WINDOW_HANDLES: std::cell::RefCell> = std::cell::RefCell::new(HashMap::new()); +} From ec184e497cade04b7e465f8e074aa4e2f6a190e7 Mon Sep 17 00:00:00 2001 From: louis030195 Date: Thu, 21 Nov 2024 00:44:40 +0000 Subject: [PATCH 6/8] remove bin --- screenpipe-vision/bin | 316 ------------------------------------------ 1 file changed, 316 deletions(-) delete mode 100644 screenpipe-vision/bin diff --git a/screenpipe-vision/bin b/screenpipe-vision/bin deleted file mode 100644 index 4083b7a70e..0000000000 --- a/screenpipe-vision/bin +++ /dev/null @@ -1,316 +0,0 @@ -#[cfg(target_os = "windows")] -use windows::Win32::System::RemoteDesktop::{ - WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, - WTS_INFO_CLASS, WTS_SESSION_INFOW, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::Graphics::Gdi::{ - BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, - GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::Foundation::HANDLE; - -#[cfg(target_os = "windows")] -use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; - -#[cfg(target_os = "windows")] -use windows::Win32::Security::{ - GetTokenInformation, LookupPrivilegeNameW, TokenPrivileges, SE_PRIVILEGE_ENABLED, - TOKEN_PRIVILEGES, TOKEN_QUERY, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::System::Memory::{LocalAlloc, LPTR}; - -#[cfg(target_os = "windows")] -use windows::core::PWSTR; - -use std::ffi::c_void; -use std::fs; -use std::path::Path; - -#[cfg(target_os = "windows")] -async fn capture_all_sessions() -> anyhow::Result<()> { - // Create screenshots directory if it doesn't exist - let screenshots_dir = Path::new("screenshots"); - if !screenshots_dir.exists() { - println!("creating screenshots directory..."); - fs::create_dir_all(screenshots_dir)?; - } - - loop { - println!("capturing new round of screenshots..."); - - let mut session_count: u32 = 0; - let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); - - unsafe { - WTSEnumerateSessionsW( - WTS_CURRENT_SERVER_HANDLE, - 0, - 1, - &mut sessions, - &mut session_count, - )?; - } - - let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - - println!("found {} sessions", session_count); - - for session in sessions_slice { - if session.State == WTSActive { - println!("processing session {}", session.SessionId); - - // Get session username to verify we have access - let mut bytes_returned: u32 = 0; - let mut username = PWSTR::null(); - - unsafe { - match WTSQuerySessionInformationW( - WTS_CURRENT_SERVER_HANDLE, - session.SessionId, - WTS_INFO_CLASS(5), // WTSUserName - &mut username, - &mut bytes_returned, - ) { - Ok(_) => { - let username_str = username.to_string()?; - println!( - "session {} belongs to user: {}", - session.SessionId, username_str - ); - } - Err(e) => { - println!( - "warning: couldn't get username for session {}: {:?}", - session.SessionId, e - ); - } - } - } - - // Create DC for this session - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() - }; - - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); - - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; - - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); - } - } else { - println!("skipping inactive session {}", session.SessionId); - } - } - - // Add delay between captures - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - } - - Ok(()) -} - -#[cfg(target_os = "windows")] -fn save_buffer_as_png( - buffer: &[u8], - width: u32, - height: u32, - filename: &str, -) -> anyhow::Result<()> { - use image::{ImageBuffer, Rgba}; - - let mut img = ImageBuffer::new(width, height); - - for (x, y, pixel) in img.enumerate_pixels_mut() { - let pos = ((y * width + x) * 4) as usize; - *pixel = Rgba([ - buffer[pos + 2], // B - buffer[pos + 1], // G - buffer[pos], // R - buffer[pos + 3], // A - ]); - } - - img.save(filename)?; - println!("saved capture to {}", filename); - Ok(()) -} - -#[cfg(target_os = "windows")] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("checking rdp permissions..."); - - unsafe { - let mut token_handle = HANDLE::default(); - OpenProcessToken( - GetCurrentProcess(), - TOKEN_QUERY, - &mut token_handle as *mut HANDLE, - )?; - println!("successfully opened process token"); - - let mut return_length = 0; - let result = - GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); - - if result.is_err() { - println!("got required buffer size: {} bytes", return_length); - - // Allocate the required buffer - let buffer = LocalAlloc(LPTR, return_length as usize)?; - let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; - - // Second call with properly sized buffer - GetTokenInformation( - token_handle, - TokenPrivileges, - Some(buffer.0 as *mut c_void), - return_length, - &mut return_length, - )?; - - let privileges = &*privileges_ptr; - println!("found {} privileges", privileges.PrivilegeCount); - - let privilege_array = std::slice::from_raw_parts( - privileges.Privileges.as_ptr(), - privileges.PrivilegeCount as usize, - ); - - for privilege in privilege_array { - // First call to get the required name length - let mut name_len = 0; - let name_result = - LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); - - // Ignore the expected error from getting the size - if name_len > 0 { - // Allocate buffer with the correct size (+1 for null terminator) - let mut name = vec![0u16; name_len as usize + 1]; - let mut final_len = name_len; - - // Second call to actually get the name - match LookupPrivilegeNameW( - None, - &privilege.Luid, - PWSTR(name.as_mut_ptr()), - &mut final_len, - ) { - Ok(_) => { - let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); - println!( - "privilege: {} (enabled: {})", - priv_name, - privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED - ); - } - Err(e) => { - println!("failed to get privilege name: {:?}", e); - } - } - } - } - } - } - - Ok(()) -} - -#[cfg(not(target_os = "windows"))] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("rdp permissions check only available on windows"); - Ok(()) -} - -#[cfg(not(target_os = "windows"))] -async fn capture_all_sessions() -> anyhow::Result<()> { - println!("rdp capture only available on windows"); - Ok(()) -} - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - println!("starting rdp example..."); - - println!("checking permissions..."); - let err = check_rdp_permissions().await; - if let Err(e) = err { - println!("error checking permissions: {:?}", e); - } - - println!("starting continuous capture..."); - let err = capture_all_sessions().await; - if let Err(e) = err { - println!("error capturing session: {:?}", e); - } - - println!("example completed successfully"); - Ok(()) -} From 89340a254ab99515817f57349045ba6c44b6a3d1 Mon Sep 17 00:00:00 2001 From: louis030195 Date: Fri, 22 Nov 2024 19:24:13 +0000 Subject: [PATCH 7/8] fix permission --- screenpipe-vision/examples/rdp.rs | 294 +++++++++++++++++++----------- screenpipe-vision/src/bin/rdp.rs | 294 +++++++++++++++++++----------- 2 files changed, 380 insertions(+), 208 deletions(-) diff --git a/screenpipe-vision/examples/rdp.rs b/screenpipe-vision/examples/rdp.rs index 4083b7a70e..564ae29e87 100644 --- a/screenpipe-vision/examples/rdp.rs +++ b/screenpipe-vision/examples/rdp.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "windows")] use windows::Win32::System::RemoteDesktop::{ WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, - WTS_INFO_CLASS, WTS_SESSION_INFOW, + WTS_INFO_CLASS, WTS_SESSION_INFOW, WTSVirtualChannelQuery, WTSFreeMemory, WTSQueryUserToken, }; #[cfg(target_os = "windows")] @@ -28,12 +28,97 @@ use windows::Win32::System::Memory::{LocalAlloc, LPTR}; #[cfg(target_os = "windows")] use windows::core::PWSTR; +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ImpersonateLoggedOnUser, RevertToSelf}; + use std::ffi::c_void; use std::fs; use std::path::Path; +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ + AdjustTokenPrivileges, LookupPrivilegeValueW, TOKEN_ADJUST_PRIVILEGES, + LUID_AND_ATTRIBUTES +}; + +#[cfg(target_os = "windows")] +fn enable_required_privileges() -> anyhow::Result<()> { + use windows::Win32::Foundation::CloseHandle; + + unsafe { + let mut token_handle = HANDLE::default(); + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, + &mut token_handle, + )?; + + // List of privileges we need + let required_privilege_names = [ + "SeTcbPrivilege", + "SeDebugPrivilege", + "SeImpersonatePrivilege", + "SeAssignPrimaryTokenPrivilege", + "SeIncreaseQuotaPrivilege", + ]; + + for privilege_name in required_privilege_names.iter() { + let mut tp = TOKEN_PRIVILEGES { + PrivilegeCount: 1, + Privileges: [LUID_AND_ATTRIBUTES { + Luid: windows::Win32::Foundation::LUID::default(), + Attributes: SE_PRIVILEGE_ENABLED, + }], + }; + + let priv_name = format!("{}\0", privilege_name); + let utf16_name = priv_name.encode_utf16().collect::>(); + + println!("attempting to enable {}...", privilege_name); + + if let Err(e) = LookupPrivilegeValueW( + None, + PWSTR(utf16_name.as_ptr() as *mut u16), + &mut tp.Privileges[0].Luid, + ) { + println!("warning: failed to lookup {}: {:?}", privilege_name, e); + continue; + } + + match AdjustTokenPrivileges( + token_handle, + false, + Some(&tp), + 0, + None, + None, + ) { + Ok(_) => { + let result = windows::Win32::Foundation::GetLastError(); + if result == windows::Win32::Foundation::ERROR_SUCCESS { + println!("successfully enabled {}", privilege_name); + } else { + println!("failed to enable {} (error: {:?})", privilege_name, result); + } + } + Err(e) => println!("failed to adjust privileges for {}: {:?}", privilege_name, e), + } + } + + CloseHandle(token_handle); + Ok(()) + } +} + #[cfg(target_os = "windows")] async fn capture_all_sessions() -> anyhow::Result<()> { + use windows::Win32::Foundation::CloseHandle; + + println!("enabling privileges..."); + if let Err(e) = enable_required_privileges() { + println!("failed to enable privileges: {:?}", e); + } + // Create screenshots directory if it doesn't exist let screenshots_dir = Path::new("screenshots"); if !screenshots_dir.exists() { @@ -44,6 +129,10 @@ async fn capture_all_sessions() -> anyhow::Result<()> { loop { println!("capturing new round of screenshots..."); + // Check privileges before attempting capture + println!("checking current privileges..."); + check_rdp_permissions().await?; + let mut session_count: u32 = 0; let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); @@ -63,110 +152,108 @@ async fn capture_all_sessions() -> anyhow::Result<()> { for session in sessions_slice { if session.State == WTSActive { - println!("processing session {}", session.SessionId); - - // Get session username to verify we have access - let mut bytes_returned: u32 = 0; - let mut username = PWSTR::null(); + println!("processing session {} (state: {:?})", session.SessionId, session.State); + // Get user token for the session + let mut user_token = HANDLE::default(); unsafe { - match WTSQuerySessionInformationW( - WTS_CURRENT_SERVER_HANDLE, - session.SessionId, - WTS_INFO_CLASS(5), // WTSUserName - &mut username, - &mut bytes_returned, - ) { + match WTSQueryUserToken(session.SessionId, &mut user_token) { Ok(_) => { - let username_str = username.to_string()?; - println!( - "session {} belongs to user: {}", - session.SessionId, username_str - ); + println!("successfully got user token for session {}", session.SessionId); + // Impersonate the user + if let Ok(_) = ImpersonateLoggedOnUser(user_token) { + // Now create DC and capture screen + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) + }; + + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + } + + // Revert impersonation when done + RevertToSelf()?; + } + CloseHandle(user_token); } Err(e) => { println!( - "warning: couldn't get username for session {}: {:?}", - session.SessionId, e + "failed to get user token for session {}: {:?}", + session.SessionId, + e ); + println!("please ensure you're running as SYSTEM using: psexec -s -i cmd.exe"); + println!("or try: runas /user:SYSTEM "); + return Err(anyhow::anyhow!("Access denied - needs to run as SYSTEM")); } } } - - // Create DC for this session - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() - }; - - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); - - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; - - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); - } } else { - println!("skipping inactive session {}", session.SessionId); + println!("skipping inactive session {} (state: {:?})", session.SessionId, session.State); } } @@ -299,18 +386,17 @@ async fn capture_all_sessions() -> anyhow::Result<()> { async fn main() -> anyhow::Result<()> { println!("starting rdp example..."); - println!("checking permissions..."); - let err = check_rdp_permissions().await; - if let Err(e) = err { - println!("error checking permissions: {:?}", e); - } + println!("checking permissions before enabling privileges..."); + check_rdp_permissions().await?; + + println!("enabling privileges..."); + enable_required_privileges()?; + + println!("checking permissions after enabling privileges..."); + check_rdp_permissions().await?; println!("starting continuous capture..."); - let err = capture_all_sessions().await; - if let Err(e) = err { - println!("error capturing session: {:?}", e); - } + capture_all_sessions().await?; - println!("example completed successfully"); Ok(()) } diff --git a/screenpipe-vision/src/bin/rdp.rs b/screenpipe-vision/src/bin/rdp.rs index 4083b7a70e..564ae29e87 100644 --- a/screenpipe-vision/src/bin/rdp.rs +++ b/screenpipe-vision/src/bin/rdp.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "windows")] use windows::Win32::System::RemoteDesktop::{ WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, - WTS_INFO_CLASS, WTS_SESSION_INFOW, + WTS_INFO_CLASS, WTS_SESSION_INFOW, WTSVirtualChannelQuery, WTSFreeMemory, WTSQueryUserToken, }; #[cfg(target_os = "windows")] @@ -28,12 +28,97 @@ use windows::Win32::System::Memory::{LocalAlloc, LPTR}; #[cfg(target_os = "windows")] use windows::core::PWSTR; +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ImpersonateLoggedOnUser, RevertToSelf}; + use std::ffi::c_void; use std::fs; use std::path::Path; +#[cfg(target_os = "windows")] +use windows::Win32::Security::{ + AdjustTokenPrivileges, LookupPrivilegeValueW, TOKEN_ADJUST_PRIVILEGES, + LUID_AND_ATTRIBUTES +}; + +#[cfg(target_os = "windows")] +fn enable_required_privileges() -> anyhow::Result<()> { + use windows::Win32::Foundation::CloseHandle; + + unsafe { + let mut token_handle = HANDLE::default(); + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, + &mut token_handle, + )?; + + // List of privileges we need + let required_privilege_names = [ + "SeTcbPrivilege", + "SeDebugPrivilege", + "SeImpersonatePrivilege", + "SeAssignPrimaryTokenPrivilege", + "SeIncreaseQuotaPrivilege", + ]; + + for privilege_name in required_privilege_names.iter() { + let mut tp = TOKEN_PRIVILEGES { + PrivilegeCount: 1, + Privileges: [LUID_AND_ATTRIBUTES { + Luid: windows::Win32::Foundation::LUID::default(), + Attributes: SE_PRIVILEGE_ENABLED, + }], + }; + + let priv_name = format!("{}\0", privilege_name); + let utf16_name = priv_name.encode_utf16().collect::>(); + + println!("attempting to enable {}...", privilege_name); + + if let Err(e) = LookupPrivilegeValueW( + None, + PWSTR(utf16_name.as_ptr() as *mut u16), + &mut tp.Privileges[0].Luid, + ) { + println!("warning: failed to lookup {}: {:?}", privilege_name, e); + continue; + } + + match AdjustTokenPrivileges( + token_handle, + false, + Some(&tp), + 0, + None, + None, + ) { + Ok(_) => { + let result = windows::Win32::Foundation::GetLastError(); + if result == windows::Win32::Foundation::ERROR_SUCCESS { + println!("successfully enabled {}", privilege_name); + } else { + println!("failed to enable {} (error: {:?})", privilege_name, result); + } + } + Err(e) => println!("failed to adjust privileges for {}: {:?}", privilege_name, e), + } + } + + CloseHandle(token_handle); + Ok(()) + } +} + #[cfg(target_os = "windows")] async fn capture_all_sessions() -> anyhow::Result<()> { + use windows::Win32::Foundation::CloseHandle; + + println!("enabling privileges..."); + if let Err(e) = enable_required_privileges() { + println!("failed to enable privileges: {:?}", e); + } + // Create screenshots directory if it doesn't exist let screenshots_dir = Path::new("screenshots"); if !screenshots_dir.exists() { @@ -44,6 +129,10 @@ async fn capture_all_sessions() -> anyhow::Result<()> { loop { println!("capturing new round of screenshots..."); + // Check privileges before attempting capture + println!("checking current privileges..."); + check_rdp_permissions().await?; + let mut session_count: u32 = 0; let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); @@ -63,110 +152,108 @@ async fn capture_all_sessions() -> anyhow::Result<()> { for session in sessions_slice { if session.State == WTSActive { - println!("processing session {}", session.SessionId); - - // Get session username to verify we have access - let mut bytes_returned: u32 = 0; - let mut username = PWSTR::null(); + println!("processing session {} (state: {:?})", session.SessionId, session.State); + // Get user token for the session + let mut user_token = HANDLE::default(); unsafe { - match WTSQuerySessionInformationW( - WTS_CURRENT_SERVER_HANDLE, - session.SessionId, - WTS_INFO_CLASS(5), // WTSUserName - &mut username, - &mut bytes_returned, - ) { + match WTSQueryUserToken(session.SessionId, &mut user_token) { Ok(_) => { - let username_str = username.to_string()?; - println!( - "session {} belongs to user: {}", - session.SessionId, username_str - ); + println!("successfully got user token for session {}", session.SessionId); + // Impersonate the user + if let Ok(_) = ImpersonateLoggedOnUser(user_token) { + // Now create DC and capture screen + let dc_name = format!("DISPLAY#{}\0", session.SessionId); + println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); + + let hdc = unsafe { + CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(dc_name.as_bytes().as_ptr()), + None, + None, + ) + }; + + // Create compatible DC and bitmap + let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; + let hbitmap = unsafe { + CreateCompatibleBitmap( + hdc, 1920, // width - you might want to get this from session info + 1080, // height - you might want to get this from session info + ) + }; + + unsafe { + SelectObject(hdc_mem, hbitmap); + BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; + + // Setup bitmap info + let mut bi = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: 1920, + biHeight: -1080, // Negative for top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + ..Default::default() + }; + + // Get the actual pixels + let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; + GetDIBits( + hdc_mem, + hbitmap, + 0, + 1080, + Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), + &mut bi, + DIB_RGB_COLORS, + ); + + // Save the image + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + + let filename = format!( + "screenshots/session_{}_capture_{}.png", + session.SessionId, timestamp + ); + save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + + // Cleanup + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(hdc); + } + + // Revert impersonation when done + RevertToSelf()?; + } + CloseHandle(user_token); } Err(e) => { println!( - "warning: couldn't get username for session {}: {:?}", - session.SessionId, e + "failed to get user token for session {}: {:?}", + session.SessionId, + e ); + println!("please ensure you're running as SYSTEM using: psexec -s -i cmd.exe"); + println!("or try: runas /user:SYSTEM "); + return Err(anyhow::anyhow!("Access denied - needs to run as SYSTEM")); } } } - - // Create DC for this session - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() - }; - - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); - - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; - - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); - } } else { - println!("skipping inactive session {}", session.SessionId); + println!("skipping inactive session {} (state: {:?})", session.SessionId, session.State); } } @@ -299,18 +386,17 @@ async fn capture_all_sessions() -> anyhow::Result<()> { async fn main() -> anyhow::Result<()> { println!("starting rdp example..."); - println!("checking permissions..."); - let err = check_rdp_permissions().await; - if let Err(e) = err { - println!("error checking permissions: {:?}", e); - } + println!("checking permissions before enabling privileges..."); + check_rdp_permissions().await?; + + println!("enabling privileges..."); + enable_required_privileges()?; + + println!("checking permissions after enabling privileges..."); + check_rdp_permissions().await?; println!("starting continuous capture..."); - let err = capture_all_sessions().await; - if let Err(e) = err { - println!("error capturing session: {:?}", e); - } + capture_all_sessions().await?; - println!("example completed successfully"); Ok(()) } From ae512096cf3918092fc6186f9900a7a1132dd687 Mon Sep 17 00:00:00 2001 From: louis030195 Date: Thu, 28 Nov 2024 16:19:42 +0000 Subject: [PATCH 8/8] progress --- screenpipe-vision/Cargo.toml | 4 +- screenpipe-vision/examples/rdp.rs | 402 --------------- screenpipe-vision/src/bin/rdp.rs | 796 ++++++++++++++++++++---------- screenpipe-vision/src/bin/rdp2.rs | 244 +++++++++ 4 files changed, 772 insertions(+), 674 deletions(-) delete mode 100644 screenpipe-vision/examples/rdp.rs create mode 100644 screenpipe-vision/src/bin/rdp2.rs diff --git a/screenpipe-vision/Cargo.toml b/screenpipe-vision/Cargo.toml index 58b25c29e4..327a9b9c94 100644 --- a/screenpipe-vision/Cargo.toml +++ b/screenpipe-vision/Cargo.toml @@ -91,8 +91,10 @@ name = "rdp" path = "examples/rdp.rs" [target.'cfg(target_os = "windows")'.dependencies] -windows = { version = "0.58", features = ["Win32_System_ProcessStatus", "Win32_UI_WindowsAndMessaging", "Win32_UI", "Win32_Graphics_Gdi", "Win32_Graphics", "Win32_System_Memory", "Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams", "Win32_System_RemoteDesktop", "Win32_System_Diagnostics", "Win32_System_Diagnostics_Debug", "Win32_Security", "Win32_System_Threading", "Win32_System_Registry", "Win32_System_StationsAndDesktops"] } +windows = { version = "0.58", features = ["Win32_System_WindowsProgramming", "Win32_System_SystemInformation", "Win32_System_ProcessStatus", "Win32_UI_WindowsAndMessaging", "Win32_UI", "Win32_Graphics_Gdi", "Win32_Graphics", "Win32_System_Memory", "Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams", "Win32_System_RemoteDesktop", "Win32_System_Diagnostics", "Win32_System_Diagnostics_Debug", "Win32_Security", "Win32_System_Threading", "Win32_System_Registry", "Win32_System_StationsAndDesktops"] } +windows-sys = { version = "0.59", features = ["Win32_UI_WindowsAndMessaging", "Win32_System_StationsAndDesktops"] } xcap = "0.0.12" +byteorder = "1.4" [target.'cfg(target_os = "macos")'.dependencies] libc = "0.2" diff --git a/screenpipe-vision/examples/rdp.rs b/screenpipe-vision/examples/rdp.rs deleted file mode 100644 index 564ae29e87..0000000000 --- a/screenpipe-vision/examples/rdp.rs +++ /dev/null @@ -1,402 +0,0 @@ -#[cfg(target_os = "windows")] -use windows::Win32::System::RemoteDesktop::{ - WTSActive, WTSEnumerateSessionsW, WTSQuerySessionInformationW, WTS_CURRENT_SERVER_HANDLE, - WTS_INFO_CLASS, WTS_SESSION_INFOW, WTSVirtualChannelQuery, WTSFreeMemory, WTSQueryUserToken, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::Graphics::Gdi::{ - BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, CreateDCA, DeleteDC, DeleteObject, - GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, SRCCOPY, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::Foundation::HANDLE; - -#[cfg(target_os = "windows")] -use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; - -#[cfg(target_os = "windows")] -use windows::Win32::Security::{ - GetTokenInformation, LookupPrivilegeNameW, TokenPrivileges, SE_PRIVILEGE_ENABLED, - TOKEN_PRIVILEGES, TOKEN_QUERY, -}; - -#[cfg(target_os = "windows")] -use windows::Win32::System::Memory::{LocalAlloc, LPTR}; - -#[cfg(target_os = "windows")] -use windows::core::PWSTR; - -#[cfg(target_os = "windows")] -use windows::Win32::Security::{ImpersonateLoggedOnUser, RevertToSelf}; - -use std::ffi::c_void; -use std::fs; -use std::path::Path; - -#[cfg(target_os = "windows")] -use windows::Win32::Security::{ - AdjustTokenPrivileges, LookupPrivilegeValueW, TOKEN_ADJUST_PRIVILEGES, - LUID_AND_ATTRIBUTES -}; - -#[cfg(target_os = "windows")] -fn enable_required_privileges() -> anyhow::Result<()> { - use windows::Win32::Foundation::CloseHandle; - - unsafe { - let mut token_handle = HANDLE::default(); - OpenProcessToken( - GetCurrentProcess(), - TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, - &mut token_handle, - )?; - - // List of privileges we need - let required_privilege_names = [ - "SeTcbPrivilege", - "SeDebugPrivilege", - "SeImpersonatePrivilege", - "SeAssignPrimaryTokenPrivilege", - "SeIncreaseQuotaPrivilege", - ]; - - for privilege_name in required_privilege_names.iter() { - let mut tp = TOKEN_PRIVILEGES { - PrivilegeCount: 1, - Privileges: [LUID_AND_ATTRIBUTES { - Luid: windows::Win32::Foundation::LUID::default(), - Attributes: SE_PRIVILEGE_ENABLED, - }], - }; - - let priv_name = format!("{}\0", privilege_name); - let utf16_name = priv_name.encode_utf16().collect::>(); - - println!("attempting to enable {}...", privilege_name); - - if let Err(e) = LookupPrivilegeValueW( - None, - PWSTR(utf16_name.as_ptr() as *mut u16), - &mut tp.Privileges[0].Luid, - ) { - println!("warning: failed to lookup {}: {:?}", privilege_name, e); - continue; - } - - match AdjustTokenPrivileges( - token_handle, - false, - Some(&tp), - 0, - None, - None, - ) { - Ok(_) => { - let result = windows::Win32::Foundation::GetLastError(); - if result == windows::Win32::Foundation::ERROR_SUCCESS { - println!("successfully enabled {}", privilege_name); - } else { - println!("failed to enable {} (error: {:?})", privilege_name, result); - } - } - Err(e) => println!("failed to adjust privileges for {}: {:?}", privilege_name, e), - } - } - - CloseHandle(token_handle); - Ok(()) - } -} - -#[cfg(target_os = "windows")] -async fn capture_all_sessions() -> anyhow::Result<()> { - use windows::Win32::Foundation::CloseHandle; - - println!("enabling privileges..."); - if let Err(e) = enable_required_privileges() { - println!("failed to enable privileges: {:?}", e); - } - - // Create screenshots directory if it doesn't exist - let screenshots_dir = Path::new("screenshots"); - if !screenshots_dir.exists() { - println!("creating screenshots directory..."); - fs::create_dir_all(screenshots_dir)?; - } - - loop { - println!("capturing new round of screenshots..."); - - // Check privileges before attempting capture - println!("checking current privileges..."); - check_rdp_permissions().await?; - - let mut session_count: u32 = 0; - let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); - - unsafe { - WTSEnumerateSessionsW( - WTS_CURRENT_SERVER_HANDLE, - 0, - 1, - &mut sessions, - &mut session_count, - )?; - } - - let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - - println!("found {} sessions", session_count); - - for session in sessions_slice { - if session.State == WTSActive { - println!("processing session {} (state: {:?})", session.SessionId, session.State); - - // Get user token for the session - let mut user_token = HANDLE::default(); - unsafe { - match WTSQueryUserToken(session.SessionId, &mut user_token) { - Ok(_) => { - println!("successfully got user token for session {}", session.SessionId); - // Impersonate the user - if let Ok(_) = ImpersonateLoggedOnUser(user_token) { - // Now create DC and capture screen - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() - }; - - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); - - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; - - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); - } - - // Revert impersonation when done - RevertToSelf()?; - } - CloseHandle(user_token); - } - Err(e) => { - println!( - "failed to get user token for session {}: {:?}", - session.SessionId, - e - ); - println!("please ensure you're running as SYSTEM using: psexec -s -i cmd.exe"); - println!("or try: runas /user:SYSTEM "); - return Err(anyhow::anyhow!("Access denied - needs to run as SYSTEM")); - } - } - } - } else { - println!("skipping inactive session {} (state: {:?})", session.SessionId, session.State); - } - } - - // Add delay between captures - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - } - - Ok(()) -} - -#[cfg(target_os = "windows")] -fn save_buffer_as_png( - buffer: &[u8], - width: u32, - height: u32, - filename: &str, -) -> anyhow::Result<()> { - use image::{ImageBuffer, Rgba}; - - let mut img = ImageBuffer::new(width, height); - - for (x, y, pixel) in img.enumerate_pixels_mut() { - let pos = ((y * width + x) * 4) as usize; - *pixel = Rgba([ - buffer[pos + 2], // B - buffer[pos + 1], // G - buffer[pos], // R - buffer[pos + 3], // A - ]); - } - - img.save(filename)?; - println!("saved capture to {}", filename); - Ok(()) -} - -#[cfg(target_os = "windows")] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("checking rdp permissions..."); - - unsafe { - let mut token_handle = HANDLE::default(); - OpenProcessToken( - GetCurrentProcess(), - TOKEN_QUERY, - &mut token_handle as *mut HANDLE, - )?; - println!("successfully opened process token"); - - let mut return_length = 0; - let result = - GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); - - if result.is_err() { - println!("got required buffer size: {} bytes", return_length); - - // Allocate the required buffer - let buffer = LocalAlloc(LPTR, return_length as usize)?; - let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; - - // Second call with properly sized buffer - GetTokenInformation( - token_handle, - TokenPrivileges, - Some(buffer.0 as *mut c_void), - return_length, - &mut return_length, - )?; - - let privileges = &*privileges_ptr; - println!("found {} privileges", privileges.PrivilegeCount); - - let privilege_array = std::slice::from_raw_parts( - privileges.Privileges.as_ptr(), - privileges.PrivilegeCount as usize, - ); - - for privilege in privilege_array { - // First call to get the required name length - let mut name_len = 0; - let name_result = - LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); - - // Ignore the expected error from getting the size - if name_len > 0 { - // Allocate buffer with the correct size (+1 for null terminator) - let mut name = vec![0u16; name_len as usize + 1]; - let mut final_len = name_len; - - // Second call to actually get the name - match LookupPrivilegeNameW( - None, - &privilege.Luid, - PWSTR(name.as_mut_ptr()), - &mut final_len, - ) { - Ok(_) => { - let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); - println!( - "privilege: {} (enabled: {})", - priv_name, - privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED - ); - } - Err(e) => { - println!("failed to get privilege name: {:?}", e); - } - } - } - } - } - } - - Ok(()) -} - -#[cfg(not(target_os = "windows"))] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("rdp permissions check only available on windows"); - Ok(()) -} - -#[cfg(not(target_os = "windows"))] -async fn capture_all_sessions() -> anyhow::Result<()> { - println!("rdp capture only available on windows"); - Ok(()) -} - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - println!("starting rdp example..."); - - println!("checking permissions before enabling privileges..."); - check_rdp_permissions().await?; - - println!("enabling privileges..."); - enable_required_privileges()?; - - println!("checking permissions after enabling privileges..."); - check_rdp_permissions().await?; - - println!("starting continuous capture..."); - capture_all_sessions().await?; - - Ok(()) -} diff --git a/screenpipe-vision/src/bin/rdp.rs b/screenpipe-vision/src/bin/rdp.rs index 564ae29e87..641f826b48 100644 --- a/screenpipe-vision/src/bin/rdp.rs +++ b/screenpipe-vision/src/bin/rdp.rs @@ -42,82 +42,134 @@ use windows::Win32::Security::{ }; #[cfg(target_os = "windows")] -fn enable_required_privileges() -> anyhow::Result<()> { - use windows::Win32::Foundation::CloseHandle; +use windows::Win32::{ + Foundation::{CloseHandle, INVALID_HANDLE_VALUE}, + Security::{ + DuplicateTokenEx, SecurityImpersonation, SecurityIdentification, + TokenPrimary, TOKEN_ALL_ACCESS, + }, +}; + + +#[cfg(target_os = "windows")] +use windows::Win32::Graphics::Gdi::{ + EnumDisplaySettingsW, DEVMODEW, ENUM_CURRENT_SETTINGS, +}; + +use std::env; + +#[cfg(target_os = "windows")] +fn check_privilege(privilege_name: &str) -> bool { + use windows::core::PCWSTR; unsafe { - let mut token_handle = HANDLE::default(); - OpenProcessToken( - GetCurrentProcess(), - TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, - &mut token_handle, - )?; + let mut token = HANDLE::default(); + let mut has_privilege = false; - // List of privileges we need - let required_privilege_names = [ - "SeTcbPrivilege", - "SeDebugPrivilege", - "SeImpersonatePrivilege", - "SeAssignPrimaryTokenPrivilege", - "SeIncreaseQuotaPrivilege", - ]; - - for privilege_name in required_privilege_names.iter() { - let mut tp = TOKEN_PRIVILEGES { - PrivilegeCount: 1, - Privileges: [LUID_AND_ATTRIBUTES { - Luid: windows::Win32::Foundation::LUID::default(), - Attributes: SE_PRIVILEGE_ENABLED, - }], - }; - - let priv_name = format!("{}\0", privilege_name); - let utf16_name = priv_name.encode_utf16().collect::>(); - - println!("attempting to enable {}...", privilege_name); + if OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + &mut token, + ).is_ok() { + let mut luid = windows::Win32::Foundation::LUID::default(); + let privilege_name_wide: Vec = format!("{}\0", privilege_name).encode_utf16().collect(); - if let Err(e) = LookupPrivilegeValueW( + if LookupPrivilegeValueW( None, - PWSTR(utf16_name.as_ptr() as *mut u16), - &mut tp.Privileges[0].Luid, - ) { - println!("warning: failed to lookup {}: {:?}", privilege_name, e); - continue; + PCWSTR(privilege_name_wide.as_ptr()), + &mut luid, + ).is_ok() { + let mut tp = TOKEN_PRIVILEGES::default(); + let mut return_length = 0u32; + + if GetTokenInformation( + token, + TokenPrivileges, + Some(&mut tp as *mut _ as *mut c_void), + std::mem::size_of::() as u32, + &mut return_length, + ).is_ok() { + has_privilege = true; + println!("privilege {} is already enabled", privilege_name); + } } + CloseHandle(token); + } + has_privilege + } +} - match AdjustTokenPrivileges( - token_handle, - false, - Some(&tp), - 0, - None, +#[cfg(target_os = "windows")] +fn enable_privilege(privilege_name: &str) -> anyhow::Result<()> { + use windows::core::PCWSTR; + + unsafe { + let mut token = HANDLE::default(); + if OpenProcessToken( + GetCurrentProcess(), + TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, + &mut token, + ).is_ok() { + let mut luid = windows::Win32::Foundation::LUID::default(); + let privilege_name_wide: Vec = format!("{}\0", privilege_name).encode_utf16().collect(); + + if LookupPrivilegeValueW( None, - ) { - Ok(_) => { - let result = windows::Win32::Foundation::GetLastError(); - if result == windows::Win32::Foundation::ERROR_SUCCESS { - println!("successfully enabled {}", privilege_name); - } else { - println!("failed to enable {} (error: {:?})", privilege_name, result); - } + PCWSTR(privilege_name_wide.as_ptr()), + &mut luid, + ).is_ok() { + let mut tp = TOKEN_PRIVILEGES { + PrivilegeCount: 1, + Privileges: [LUID_AND_ATTRIBUTES { + Luid: luid, + Attributes: SE_PRIVILEGE_ENABLED, + }], + }; + + if AdjustTokenPrivileges( + token, + false, + Some(&tp), + 0, + None, + None, + ).is_ok() { + println!("successfully enabled privilege: {}", privilege_name); + } else { + println!("failed to adjust token privileges: {:?}", windows::core::Error::from_win32()); } - Err(e) => println!("failed to adjust privileges for {}: {:?}", privilege_name, e), + } else { + println!("failed to lookup privilege value: {:?}", windows::core::Error::from_win32()); } + CloseHandle(token); } - - CloseHandle(token_handle); - Ok(()) } + Ok(()) } + #[cfg(target_os = "windows")] async fn capture_all_sessions() -> anyhow::Result<()> { - use windows::Win32::Foundation::CloseHandle; - - println!("enabling privileges..."); - if let Err(e) = enable_required_privileges() { - println!("failed to enable privileges: {:?}", e); + // Only enable privileges if we don't already have them + let required_privileges = [ + "SeSecurityPrivilege", + "SeBackupPrivilege", + "SeRestorePrivilege", + "SeImpersonatePrivilege", + "SeAssignPrimaryTokenPrivilege", + "SeTcbPrivilege", + "SeIncreaseQuotaPrivilege" + ]; + + for privilege in required_privileges { + if !check_privilege(privilege) { + enable_privilege(privilege)?; + } } + + use windows::Win32::System::{RemoteDesktop::{WTSCloseServer, WTSOpenServerW, WTS_CONNECTSTATE_CLASS}, WindowsProgramming::GetComputerNameW}; + + // Create screenshots directory if it doesn't exist let screenshots_dir = Path::new("screenshots"); @@ -126,144 +178,439 @@ async fn capture_all_sessions() -> anyhow::Result<()> { fs::create_dir_all(screenshots_dir)?; } - loop { - println!("capturing new round of screenshots..."); - - // Check privileges before attempting capture - println!("checking current privileges..."); - check_rdp_permissions().await?; - - let mut session_count: u32 = 0; - let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + // Get the computer name + let mut computer_name: Vec = vec![0; 256]; + let mut size = computer_name.len() as u32; + + unsafe { + if GetComputerNameW(PWSTR(computer_name.as_mut_ptr()), &mut size).is_err() { + println!("Failed to get computer name: {:?}", windows::core::Error::from_win32()); + } + } - unsafe { - WTSEnumerateSessionsW( - WTS_CURRENT_SERVER_HANDLE, - 0, - 1, - &mut sessions, - &mut session_count, - )?; + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + + let server_handle = unsafe { + let server_url = env::var("RDP_SERVER_URL") + .unwrap_or_else(|_| { + println!("warning: RDP_SERVER_URL not set, using localhost"); + "localhost".to_string() + }); + + let server_url_wide: Vec = format!("{}\0", server_url).encode_utf16().collect(); + let server_handle = WTSOpenServerW(PWSTR(server_url_wide.as_ptr() as *mut u16)); + // ... existing code ... + if server_handle.is_invalid() { + println!("Failed to open server: {:?}", windows::core::Error::from_win32()); + return Ok(()); } - let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - - println!("found {} sessions", session_count); - - for session in sessions_slice { - if session.State == WTSActive { - println!("processing session {} (state: {:?})", session.SessionId, session.State); - - // Get user token for the session - let mut user_token = HANDLE::default(); - unsafe { - match WTSQueryUserToken(session.SessionId, &mut user_token) { - Ok(_) => { - println!("successfully got user token for session {}", session.SessionId); - // Impersonate the user - if let Ok(_) = ImpersonateLoggedOnUser(user_token) { - // Now create DC and capture screen - let dc_name = format!("DISPLAY#{}\0", session.SessionId); - println!("creating DC with name: {}", dc_name.trim_end_matches('\0')); - - let hdc = unsafe { - CreateDCA( - windows::core::PCSTR(b"DISPLAY\0".as_ptr()), - windows::core::PCSTR(dc_name.as_bytes().as_ptr()), - None, - None, - ) - }; - - // Create compatible DC and bitmap - let hdc_mem = unsafe { CreateCompatibleDC(hdc) }; - let hbitmap = unsafe { - CreateCompatibleBitmap( - hdc, 1920, // width - you might want to get this from session info - 1080, // height - you might want to get this from session info - ) - }; - - unsafe { - SelectObject(hdc_mem, hbitmap); - BitBlt(hdc_mem, 0, 0, 1920, 1080, hdc, 0, 0, SRCCOPY)?; - - // Setup bitmap info - let mut bi = BITMAPINFO { - bmiHeader: BITMAPINFOHEADER { - biSize: std::mem::size_of::() as u32, - biWidth: 1920, - biHeight: -1080, // Negative for top-down - biPlanes: 1, - biBitCount: 32, - biCompression: 0, - biSizeImage: 0, - biXPelsPerMeter: 0, - biYPelsPerMeter: 0, - biClrUsed: 0, - biClrImportant: 0, - }, - ..Default::default() - }; - - // Get the actual pixels - let mut buffer = vec![0u8; (1920 * 1080 * 4) as usize]; - GetDIBits( - hdc_mem, - hbitmap, - 0, - 1080, - Some(buffer.as_mut_ptr() as *mut std::ffi::c_void), - &mut bi, - DIB_RGB_COLORS, - ); + WTSEnumerateSessionsW( + server_handle, + 0, + 1, + &mut sessions, + &mut session_count, + )?; - // Save the image - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); + // Don't forget to close the server handle + // WTSCloseServer(server_handle); + server_handle + }; - let filename = format!( - "screenshots/session_{}_capture_{}.png", - session.SessionId, timestamp - ); - save_buffer_as_png(&buffer, 1920, 1080, &filename)?; + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; - // Cleanup - DeleteObject(hbitmap); - DeleteDC(hdc_mem); - DeleteDC(hdc); - } + println!("found {} sessions", session_count); - // Revert impersonation when done - RevertToSelf()?; - } - CloseHandle(user_token); - } - Err(e) => { - println!( - "failed to get user token for session {}: {:?}", + for session in sessions_slice { + // In the sessions loop, before checking winstation + println!("Session {}: State={} ({:?})", + session.SessionId, + session.State.0, + session.State + ); + let mut username_ptr: *mut u16 = std::ptr::null_mut(); + let mut winstation_ptr: *mut u16 = std::ptr::null_mut(); + let mut client_name_ptr: *mut u16 = std::ptr::null_mut(); + let mut bytes_returned: u32 = 0; + + unsafe { + // Get WinStation name to check if it's RDP + match WTSQuerySessionInformationW( + server_handle, + session.SessionId, + WTS_INFO_CLASS(6), + &mut winstation_ptr as *mut *mut u16 as *mut _, + &mut bytes_returned, + ) { + Ok(_) => { + let winstation = if !winstation_ptr.is_null() { + String::from_utf16_lossy(std::slice::from_raw_parts( + winstation_ptr, + (bytes_returned / 2 - 1) as usize, + )) + } else { + "Unknown".to_string() + }; + + println!("winstation: {}", winstation); + // Modified RDP session check to include more variants + if winstation.contains("RDP") || winstation.contains("Tcp") { + println!("Found RDP session {}: {}", session.SessionId, winstation); + + // Get username + if let Ok(_) = WTSQuerySessionInformationW( + server_handle, + session.SessionId, + WTS_INFO_CLASS(5), + &mut username_ptr as *mut *mut u16 as *mut _, + &mut bytes_returned, + ) { + let username = if !username_ptr.is_null() { + String::from_utf16_lossy(std::slice::from_raw_parts( + username_ptr, + (bytes_returned / 2 - 1) as usize, + )) + } else { + "Unknown".to_string() + }; + + println!("RDP Session {}: username={}, state={:?}", session.SessionId, - e + username, + session.State ); - println!("please ensure you're running as SYSTEM using: psexec -s -i cmd.exe"); - println!("or try: runas /user:SYSTEM "); - return Err(anyhow::anyhow!("Access denied - needs to run as SYSTEM")); + + // Modified session state check - accept more states + // WTS_CONNECTSTATE_CLASS values: + // 0 = Active + // 1 = Connected + // 2 = ConnectQuery + // 3 = Shadow + // 4 = Disconnected + // 5 = Idle + // 6 = Listen + // 7 = Reset + // 8 = Down + // 9 = Init + if session.State.0 > 7 { // Only skip truly inactive states + println!("Skipping inactive session {} (state={})", session.SessionId, session.State.0); + continue; + } + + // Skip if no username (system or listener sessions) + if username.is_empty() { + println!("Skipping session {} with no username", session.SessionId); + continue; + } + + + + // Try to get user token for the session + let mut token = HANDLE::default(); + if let Err(e) = WTSQueryUserToken(session.SessionId, &mut token) { + println!("Failed to get token directly for session {}, trying alternative method...", session.SessionId); + + unsafe { + let mut process_token = HANDLE::default(); + if OpenProcessToken( + GetCurrentProcess(), + TOKEN_ALL_ACCESS, + &mut process_token, + ).is_ok() { + println!("Successfully opened process token with full access"); + + // Create a new token via duplication + let mut duplicated_token = HANDLE::default(); + if DuplicateTokenEx( + process_token, + TOKEN_ALL_ACCESS, + None, + SecurityImpersonation, + TokenPrimary, + &mut duplicated_token, + ).is_ok() { + println!("Successfully duplicated token"); + token = duplicated_token; + } else { + println!("Failed to duplicate token: {:?}", windows::core::Error::from_win32()); + CloseHandle(process_token); + continue; + } + CloseHandle(process_token); + } + } + } + + // Validate token before using it + if token.is_invalid() || token == INVALID_HANDLE_VALUE { + println!("Invalid token for session {}", session.SessionId); + continue; + } + + // Now try to impersonate + if unsafe { ImpersonateLoggedOnUser(token) }.is_ok() { + println!("Successfully impersonated user for session {}", session.SessionId); + + + + if let Err(e) = capture_session_screen(session.SessionId, &screenshots_dir) { + println!("Failed to capture RDP session: {:?}", e); + } + + // Make sure to revert when done + unsafe { RevertToSelf() }; + } else { + println!( + "Failed to impersonate user for session {}: {:?}", + session.SessionId, + windows::core::Error::from_win32() + ); + } + // Always clean up the token + unsafe { CloseHandle(token) }; } + } else { + println!("Skipping non-RDP session {}: {}", session.SessionId, winstation); } } - } else { - println!("skipping inactive session {} (state: {:?})", session.SessionId, session.State); + Err(e) => println!("Failed to get winstation name: {:?}", e), } - } - // Add delay between captures - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + // Cleanup + if !username_ptr.is_null() { WTSFreeMemory(username_ptr as *mut _); } + if !winstation_ptr.is_null() { WTSFreeMemory(winstation_ptr as *mut _); } + if !client_name_ptr.is_null() { WTSFreeMemory(client_name_ptr as *mut _); } + } + } + unsafe { + WTSCloseServer(server_handle); } + // Add delay between captures + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + Ok(()) } +use windows::Win32::Graphics::Gdi::HDC; +use windows::Win32::UI::WindowsAndMessaging::GetSystemMetrics; + +#[cfg(target_os = "windows")] +fn capture_session_screen(session_id: u32, screenshots_dir: &Path) -> anyhow::Result<()> { + use windows::{core::PCWSTR, Win32::Graphics::Gdi::RGBQUAD}; + + unsafe { + // First try to enumerate the actual RDP display devices + let mut dev_num = 0; + let mut display_device = windows::Win32::Graphics::Gdi::DISPLAY_DEVICEW::default(); + display_device.cb = std::mem::size_of::() as u32; + + let mut rdp_display = None; + + // First enumerate primary display devices + while windows::Win32::Graphics::Gdi::EnumDisplayDevicesW( + None, // NULL for primary devices + dev_num, + &mut display_device, + 0 // Don't get child devices yet + ).as_bool() { + // For each primary device, check its children for RDP displays + let mut child_num = 0; + let mut child_device = windows::Win32::Graphics::Gdi::DISPLAY_DEVICEW::default(); + child_device.cb = std::mem::size_of::() as u32; + println!("display_device: {:?}", display_device.DeviceName); + + while windows::Win32::Graphics::Gdi::EnumDisplayDevicesW( + PCWSTR(display_device.DeviceName.as_ptr()), + child_num, + &mut child_device, + 1 + ).as_bool() { + // Convert device string to readable format for logging + let device_string = String::from_utf16_lossy( + &child_device.DeviceString[..].iter() + .take_while(|&&c| c != 0) + .map(|&c| c) + .collect::>() + ); + + println!("checking device: {}", device_string); + + // Check if this is an RDP display + if device_string.contains("RDPDD") || device_string.contains("RDP") { + // Convert device name to string + let device_name = String::from_utf16_lossy( + &child_device.DeviceName[..].iter() + .take_while(|&&c| c != 0) + .map(|&c| c) + .collect::>() + ); + + // Check if device is active + if (child_device.StateFlags & 0x1) != 0 { // DISPLAY_DEVICE_ACTIVE + rdp_display = Some(device_name.clone()); + println!("found active RDP display: {}", device_name); + break; + } + } + + child_num += 1; + } + + if rdp_display.is_some() { + break; + } + dev_num += 1; + } + + let display_name = if let Some(name) = rdp_display { + name + } else { + // Try fallback display name format + format!("\\\\.\\DISPLAY1") // Use local device format + }; + + println!("using display name: {}", display_name); + + // Create DC directly from the display device + let mut dev_mode = DEVMODEW::default(); + dev_mode.dmSize = std::mem::size_of::() as u16; + + // Convert display name to wide string + let display_name_w: Vec = format!("{}\0", display_name).encode_utf16().collect(); + + if !EnumDisplaySettingsW( + PCWSTR(display_name_w.as_ptr()), + ENUM_CURRENT_SETTINGS, + &mut dev_mode, + ).as_bool() { + println!("failed to get display settings for {}: {:?}", + display_name, + windows::core::Error::from_win32()); + return Ok(()); + } + + let width = dev_mode.dmPelsWidth as i32; + let height = dev_mode.dmPelsHeight as i32; + + println!("found RDP display settings: {}x{} for {}", width, height, display_name); + + // Create a new DC for the remote display + let remote_dc = CreateDCA( + windows::core::PCSTR(b"DISPLAY\0".as_ptr()), + windows::core::PCSTR(display_name.as_bytes().as_ptr()), + None, + None, + ); + + if remote_dc.is_invalid() { + println!("failed to create remote DC: {:?}", windows::core::Error::from_win32()); + return Ok(()); + } + + // Rest of the capture code using remote_dc instead of hdc + let hdc_mem = CreateCompatibleDC(remote_dc); + if hdc_mem.is_invalid() { + DeleteDC(remote_dc); + println!("failed to create compatible DC: {:?}", windows::core::Error::from_win32()); + return Ok(()); + } + + // ... rest of the code remains the same but use remote_dc instead of hdc ... + let hbitmap = CreateCompatibleBitmap(remote_dc, width, height); + if hbitmap.is_invalid() { + DeleteDC(hdc_mem); + DeleteDC(remote_dc); + println!("failed to create compatible bitmap: {:?}", windows::core::Error::from_win32()); + return Ok(()); + } + + // Select bitmap into DC + let old_obj = SelectObject(hdc_mem, hbitmap); + + // Perform the bit block transfer + if let Err(e) = BitBlt( + hdc_mem, + 0, + 0, + width, + height, + remote_dc, + 0, + 0, + SRCCOPY, + ) { + SelectObject(hdc_mem, old_obj); + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + DeleteDC(remote_dc); + println!("failed to perform BitBlt: {:?}", e); + return Ok(()); + } + + // Prepare buffer for the image data + let mut buffer = vec![0u8; (width * height * 4) as usize]; + let mut info = BITMAPINFO { + bmiHeader: BITMAPINFOHEADER { + biSize: std::mem::size_of::() as u32, + biWidth: width, + biHeight: -height, // Top-down + biPlanes: 1, + biBitCount: 32, + biCompression: 0, + biSizeImage: 0, + biXPelsPerMeter: 0, + biYPelsPerMeter: 0, + biClrUsed: 0, + biClrImportant: 0, + }, + bmiColors: [RGBQUAD::default()], + }; + + // Get the bits from the bitmap + let scan_lines = GetDIBits( + hdc_mem, + hbitmap, + 0, + height as u32, + Some(buffer.as_mut_ptr() as *mut c_void), + &mut info, + DIB_RGB_COLORS, + ); + + if scan_lines == 0 { + println!("failed to get DIB bits: {:?}", windows::core::Error::from_win32()); + } else { + // Create timestamp for filename + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + let filename = screenshots_dir.join(format!("session_{}_{}.png", session_id, timestamp)); + + // Save the buffer as PNG + if let Err(e) = save_buffer_as_png(&buffer, width as u32, height as u32, filename.to_str().unwrap()) { + println!("failed to save PNG: {:?}", e); + } else { + println!("successfully saved capture for session {}", session_id); + } + } + + // Cleanup + SelectObject(hdc_mem, old_obj); + DeleteObject(hbitmap); + DeleteDC(hdc_mem); + + // Additional cleanup + DeleteDC(remote_dc); + + Ok(()) + } +} + #[cfg(target_os = "windows")] fn save_buffer_as_png( buffer: &[u8], @@ -290,91 +637,7 @@ fn save_buffer_as_png( Ok(()) } -#[cfg(target_os = "windows")] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("checking rdp permissions..."); - - unsafe { - let mut token_handle = HANDLE::default(); - OpenProcessToken( - GetCurrentProcess(), - TOKEN_QUERY, - &mut token_handle as *mut HANDLE, - )?; - println!("successfully opened process token"); - - let mut return_length = 0; - let result = - GetTokenInformation(token_handle, TokenPrivileges, None, 0, &mut return_length); - - if result.is_err() { - println!("got required buffer size: {} bytes", return_length); - - // Allocate the required buffer - let buffer = LocalAlloc(LPTR, return_length as usize)?; - let privileges_ptr = buffer.0 as *mut TOKEN_PRIVILEGES; - - // Second call with properly sized buffer - GetTokenInformation( - token_handle, - TokenPrivileges, - Some(buffer.0 as *mut c_void), - return_length, - &mut return_length, - )?; - - let privileges = &*privileges_ptr; - println!("found {} privileges", privileges.PrivilegeCount); - - let privilege_array = std::slice::from_raw_parts( - privileges.Privileges.as_ptr(), - privileges.PrivilegeCount as usize, - ); - - for privilege in privilege_array { - // First call to get the required name length - let mut name_len = 0; - let name_result = - LookupPrivilegeNameW(None, &privilege.Luid, PWSTR::null(), &mut name_len); - - // Ignore the expected error from getting the size - if name_len > 0 { - // Allocate buffer with the correct size (+1 for null terminator) - let mut name = vec![0u16; name_len as usize + 1]; - let mut final_len = name_len; - - // Second call to actually get the name - match LookupPrivilegeNameW( - None, - &privilege.Luid, - PWSTR(name.as_mut_ptr()), - &mut final_len, - ) { - Ok(_) => { - let priv_name = String::from_utf16_lossy(&name[..final_len as usize]); - println!( - "privilege: {} (enabled: {})", - priv_name, - privilege.Attributes & SE_PRIVILEGE_ENABLED == SE_PRIVILEGE_ENABLED - ); - } - Err(e) => { - println!("failed to get privilege name: {:?}", e); - } - } - } - } - } - } - - Ok(()) -} -#[cfg(not(target_os = "windows"))] -async fn check_rdp_permissions() -> anyhow::Result<()> { - println!("rdp permissions check only available on windows"); - Ok(()) -} #[cfg(not(target_os = "windows"))] async fn capture_all_sessions() -> anyhow::Result<()> { @@ -386,15 +649,6 @@ async fn capture_all_sessions() -> anyhow::Result<()> { async fn main() -> anyhow::Result<()> { println!("starting rdp example..."); - println!("checking permissions before enabling privileges..."); - check_rdp_permissions().await?; - - println!("enabling privileges..."); - enable_required_privileges()?; - - println!("checking permissions after enabling privileges..."); - check_rdp_permissions().await?; - println!("starting continuous capture..."); capture_all_sessions().await?; diff --git a/screenpipe-vision/src/bin/rdp2.rs b/screenpipe-vision/src/bin/rdp2.rs new file mode 100644 index 0000000000..82fb3c2139 --- /dev/null +++ b/screenpipe-vision/src/bin/rdp2.rs @@ -0,0 +1,244 @@ +use std::time::{SystemTime, UNIX_EPOCH}; +use byteorder::{LittleEndian, ReadBytesExt}; +use std::io::Cursor; +use std::env; + +use windows::core::{PCSTR, PWSTR}; +use windows::Win32::Foundation::*; +use windows::Win32::Graphics::Gdi::*; +use windows::Win32::System::RemoteDesktop::*; +use windows_sys::Win32::System::StationsAndDesktops::*; +use windows::Win32::UI::WindowsAndMessaging::*; +use image::{ImageBuffer, Rgba}; +use tokio::runtime::Runtime; +use windows::Win32::System::RemoteDesktop::{ + WTSVirtualChannelOpenEx +}; +use windows::core::Error; + +fn capture_remote_screen(server_handle: HANDLE, session_id: u32) { + unsafe { + // First check if the session is active + let mut session_state: u32 = 0; + let mut bytes_returned: u32 = 0; + + // if !WTSQuerySessionInformationW( + // server_handle, + // session_id, + // WTS_INFO_CLASS(5), // WTSConnectState + // &mut session_state as *mut u32 as *mut _, + // &mut bytes_returned, + // ).is_ok() { + // eprintln!("Failed to query session state: {:?}", Error::from_win32()); + // return; + // } + + // // Only proceed if session is active (WTSActive = 0) + // if session_state != 0 { + // println!("Session {} is not in active state", session_id); + // return; + // } + + // Try to open with elevated privileges + let channel = WTSVirtualChannelOpenEx( + session_id, + PCSTR("RDPGFX\0".as_ptr()), + WTS_CHANNEL_OPTION_DYNAMIC_PRI_HIGH | WTS_CHANNEL_OPTION_DYNAMIC_PRI_REAL + ); + + if channel.is_err() { + let error = Error::from_win32(); + eprintln!("Failed to open virtual channel: {:?}", error); + + // Try alternative channel names + let channel = WTSVirtualChannelOpenEx( + session_id, + PCSTR("RDPDISPLAY\0".as_ptr()), + WTS_CHANNEL_OPTION_DYNAMIC_PRI_HIGH + ); + + if channel.is_err() { + eprintln!("Failed to open RDPDISPLAY channel: {:?}", Error::from_win32()); + return; + } + } + + // Query display info + let mut buffer_size: u32 = 0; + let mut buffer: *mut std::ffi::c_void = std::ptr::null_mut(); + + if WTSVirtualChannelQuery( + channel.clone().unwrap(), + WTS_VIRTUAL_CLASS(0), + &mut buffer, + &mut buffer_size + ).is_ok() { + let buffer_slice = std::slice::from_raw_parts( + buffer as *const u8, + buffer_size as usize + ); + + // Create a cursor to read the buffer + let mut cursor = Cursor::new(buffer_slice); + + // Read header information + let header_size = cursor.read_u32::().unwrap(); + let width = cursor.read_u32::().unwrap(); + let height = cursor.read_u32::().unwrap(); + let bpp = cursor.read_u32::().unwrap(); + + println!("Remote display: {}x{} with {} bpp", width, height, bpp); + + // Calculate image data size + let stride = ((width * bpp + 31) / 32) * 4; // 32-bit aligned + let image_size = (stride * height) as usize; + + // Read the actual image data + let mut image_data = vec![0u8; image_size]; + let start_pos = header_size as usize; + image_data.copy_from_slice(&buffer_slice[start_pos..start_pos + image_size]); + + // Convert to RGBA format if necessary + let rgba_data = match bpp { + 32 => image_data, + 24 => { + let mut rgba = Vec::with_capacity(width as usize * height as usize * 4); + for chunk in image_data.chunks(3) { + rgba.extend_from_slice(chunk); + rgba.push(255); // Alpha channel + } + rgba + }, + 16 => { + let mut rgba = Vec::with_capacity(width as usize * height as usize * 4); + for chunk in image_data.chunks(2) { + let pixel = u16::from_le_bytes([chunk[0], chunk[1]]); + // 16-bit RGB565 to RGBA8888 + let r = ((pixel >> 11) & 0x1F) as u8 * 8; + let g = ((pixel >> 5) & 0x3F) as u8 * 4; + let b = (pixel & 0x1F) as u8 * 8; + rgba.extend_from_slice(&[r, g, b, 255]); + } + rgba + }, + _ => { + eprintln!("Unsupported bits per pixel: {}", bpp); + return; + } + }; + + // Create image buffer + let img_buffer = ImageBuffer::, _>::from_raw( + width, + height, + rgba_data, + ).unwrap(); + + // Save the image + let time_now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + if let Err(e) = img_buffer.save(format!("rdp_screenshot_{}.png", time_now)) { + eprintln!("failed to save screenshot: {}", e); + } else { + println!("screenshot saved to rdp_screenshot_{}.png", time_now); + } + + WTSFreeMemory(buffer); + } + + WTSVirtualChannelClose(channel.unwrap()); + } +} + +fn main() -> windows::core::Result<()> { + + + // Create Tokio runtime + let rt = Runtime::new().unwrap(); + let _guard = rt.enter(); + + unsafe { + // Enumerate all sessions on the current server + let mut session_count: u32 = 0; + let mut sessions: *mut WTS_SESSION_INFOW = std::ptr::null_mut(); + + let server_handle = unsafe { + let server_url = env::var("RDP_SERVER_URL") + .unwrap_or_else(|_| { + println!("warning: RDP_SERVER_URL not set, using localhost"); + "localhost".to_string() + }); + + let server_url_wide: Vec = format!("{}\0", server_url).encode_utf16().collect(); + let server_handle = WTSOpenServerW(PWSTR(server_url_wide.as_ptr() as *mut u16)); + // ... existing code ... + if server_handle.is_invalid() { + println!("Failed to open server: {:?}", windows::core::Error::from_win32()); + return Ok(()); + } + + WTSEnumerateSessionsW( + server_handle, + 0, + 1, + &mut sessions, + &mut session_count, + )?; + + // Don't forget to close the server handle + // WTSCloseServer(server_handle); + server_handle + }; + + let sessions_slice = unsafe { std::slice::from_raw_parts(sessions, session_count as usize) }; + + for session_info in sessions_slice { + println!("Session ID: {}", session_info.SessionId); + + // Add more session state information + let mut buffer_size: u32 = 0; + let mut buffer_ptr: *mut PWSTR = std::ptr::null_mut(); + + + + let mut winstation_ptr: *mut u16 = std::ptr::null_mut(); + let mut bytes_returned: u32 = 0; + + if WTSQuerySessionInformationW( + server_handle, + session_info.SessionId, + WTS_INFO_CLASS(6), // WTSWinStationName + &mut winstation_ptr as *mut *mut u16 as *mut _, + &mut bytes_returned, + ).is_ok() { + let winstation = if !winstation_ptr.is_null() { + String::from_utf16_lossy(std::slice::from_raw_parts( + winstation_ptr, + (bytes_returned / 2 - 1) as usize, + )) + } else { + "Unknown".to_string() + }; + + println!("Session {} winstation: {}", session_info.SessionId, winstation); + + + // Check if it's an RDP session + if winstation.contains("RDP") || winstation.contains("Tcp") || session_info.State.0 < 7 { + println!("Found RDP session: {}", session_info.SessionId); + capture_remote_screen(server_handle, session_info.SessionId); + } + + WTSFreeMemory(winstation_ptr as _); + } + } + + // Free the allocated memory + WTSFreeMemory(sessions as _); + } + + Ok(()) +}