From c3333ed2df7f7ddd6e6d9cd4d4d924964f01634f Mon Sep 17 00:00:00 2001 From: ValorBao <139973913+ValorBao@users.noreply.github.com> Date: Tue, 9 Sep 2025 23:43:16 +0800 Subject: [PATCH] Feature/screenshot (#26) * "dvt: add screen shot and change read_message in multiple message fragments" * "dvt: Add processing multiple fragments to the message module in dvt " * "cargo fmt" * Rename screen_shot to screenshot --------- Co-authored-by: Jackson Coxson --- idevice/src/services/dvt/message.rs | 58 +++++++++------- idevice/src/services/dvt/mod.rs | 1 + idevice/src/services/dvt/screenshot.rs | 64 ++++++++++++++++++ tools/Cargo.toml | 5 ++ tools/src/screenshot.rs | 94 ++++++++++++++++++++++++++ 5 files changed, 198 insertions(+), 24 deletions(-) create mode 100644 idevice/src/services/dvt/screenshot.rs create mode 100644 tools/src/screenshot.rs diff --git a/idevice/src/services/dvt/message.rs b/idevice/src/services/dvt/message.rs index 4298e77..36e4fbe 100644 --- a/idevice/src/services/dvt/message.rs +++ b/idevice/src/services/dvt/message.rs @@ -392,24 +392,35 @@ impl Message { /// # Errors /// * Various IdeviceError variants for IO and parsing failures pub async fn from_reader(reader: &mut R) -> Result { - let mut buf = [0u8; 32]; - reader.read_exact(&mut buf).await?; - - let mheader = MessageHeader { - magic: u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]), - header_len: u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]), - fragment_id: u16::from_le_bytes([buf[8], buf[9]]), - fragment_count: u16::from_le_bytes([buf[10], buf[11]]), - length: u32::from_le_bytes([buf[12], buf[13], buf[14], buf[15]]), - identifier: u32::from_le_bytes([buf[16], buf[17], buf[18], buf[19]]), - conversation_index: u32::from_le_bytes([buf[20], buf[21], buf[22], buf[23]]), - channel: u32::from_le_bytes([buf[24], buf[25], buf[26], buf[27]]), - expects_reply: u32::from_le_bytes([buf[28], buf[29], buf[30], buf[31]]) == 1, + let mut packet_data: Vec = Vec::new(); + // loop for deal with multiple fragments + let mheader = loop { + let mut buf = [0u8; 32]; + reader.read_exact(&mut buf).await?; + let header = MessageHeader { + magic: u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]), + header_len: u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]), + fragment_id: u16::from_le_bytes([buf[8], buf[9]]), + fragment_count: u16::from_le_bytes([buf[10], buf[11]]), + length: u32::from_le_bytes([buf[12], buf[13], buf[14], buf[15]]), + identifier: u32::from_le_bytes([buf[16], buf[17], buf[18], buf[19]]), + conversation_index: u32::from_le_bytes([buf[20], buf[21], buf[22], buf[23]]), + channel: u32::from_le_bytes([buf[24], buf[25], buf[26], buf[27]]), + expects_reply: u32::from_le_bytes([buf[28], buf[29], buf[30], buf[31]]) == 1, + }; + if header.fragment_count > 1 && header.fragment_id == 0 { + // when reading multiple message fragments, the first fragment contains only a message header. + continue; + } + let mut buf = vec![0u8; header.length as usize]; + reader.read_exact(&mut buf).await?; + packet_data.extend(buf); + if header.fragment_id == header.fragment_count - 1 { + break header; + } }; - - let mut buf = [0u8; 16]; - reader.read_exact(&mut buf).await?; - + // read the payload header + let buf = &packet_data[0..16]; let pheader = PayloadHeader { flags: u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]), aux_length: u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]), @@ -417,18 +428,17 @@ impl Message { buf[8], buf[9], buf[10], buf[11], buf[12], buf[13], buf[14], buf[15], ]), }; - let aux = if pheader.aux_length > 0 { - let mut buf = vec![0u8; pheader.aux_length as usize]; - reader.read_exact(&mut buf).await?; + let buf = packet_data[16..(16 + pheader.aux_length as usize)].to_vec(); Some(Aux::from_bytes(buf)?) } else { None }; - - let mut buf = vec![0u8; (pheader.total_length - pheader.aux_length as u64) as usize]; - reader.read_exact(&mut buf).await?; - + // read the data + let need_len = (pheader.total_length - pheader.aux_length as u64) as usize; + let buf = packet_data + [(pheader.aux_length + 16) as usize..pheader.aux_length as usize + 16 + need_len] + .to_vec(); let data = if buf.is_empty() { None } else { diff --git a/idevice/src/services/dvt/mod.rs b/idevice/src/services/dvt/mod.rs index d7db667..86a5fa6 100644 --- a/idevice/src/services/dvt/mod.rs +++ b/idevice/src/services/dvt/mod.rs @@ -9,6 +9,7 @@ pub mod location_simulation; pub mod message; pub mod process_control; pub mod remote_server; +pub mod screenshot; impl RsdService for remote_server::RemoteServerClient> { fn rsd_service_name() -> std::borrow::Cow<'static, str> { diff --git a/idevice/src/services/dvt/screenshot.rs b/idevice/src/services/dvt/screenshot.rs new file mode 100644 index 0000000..cd021f8 --- /dev/null +++ b/idevice/src/services/dvt/screenshot.rs @@ -0,0 +1,64 @@ +//! Screenshot service client for iOS instruments protocol. +//! +//! This module provides a client for interacting with the screenshot service +//! on iOS devices through the instruments protocol. It allows taking screenshots from the device. +//! + +use plist::Value; + +use crate::{ + IdeviceError, ReadWrite, + dvt::remote_server::{Channel, RemoteServerClient}, + obf, +}; + +/// Client for take screenshot operations on iOS devices +/// +/// Provides methods for take screnn_shot through the +/// instruments protocol. Each instance maintains its own communication channel. +pub struct ScreenshotClient<'a, R: ReadWrite> { + /// The underlying channel for communication + channel: Channel<'a, R>, +} + +impl<'a, R: ReadWrite> ScreenshotClient<'a, R> { + /// Creates a new ScreenshotClient + /// + /// # Arguments + /// * `client` - The base RemoteServerClient to use + /// + /// # Returns + /// * `Ok(ScreenshotClient)` - Connected client instance + /// * `Err(IdeviceError)` - If channel creation fails + /// + /// # Errors + /// * Propagates errors from channel creation + pub async fn new(client: &'a mut RemoteServerClient) -> Result { + let channel = client + .make_channel(obf!("com.apple.instruments.server.services.screenshot")) + .await?; // Drop `&mut client` before continuing + + Ok(Self { channel }) + } + + /// Take screenshot from the device + /// + /// # Returns + /// * `Ok(Vec)` - the bytes of the screenshot + /// * `Err(IdeviceError)` - If communication fails + /// + /// # Errors + /// * `IdeviceError::UnexpectedResponse` if server response is invalid + /// * Other communication or serialization errors + pub async fn take_screenshot(&mut self) -> Result, IdeviceError> { + let method = Value::String("takeScreenshot".into()); + + self.channel.call_method(Some(method), None, true).await?; + + let msg = self.channel.read_message().await?; + match msg.data { + Some(Value::Data(data)) => Ok(data), + _ => Err(IdeviceError::UnexpectedResponse), + } + } +} diff --git a/tools/Cargo.toml b/tools/Cargo.toml index 57b5828..3c63429 100644 --- a/tools/Cargo.toml +++ b/tools/Cargo.toml @@ -125,6 +125,11 @@ path = "src/pcapd.rs" name = "preboard" path = "src/preboard.rs" + +[[bin]] +name = "screenshot" +path = "src/screenshot.rs" + [dependencies] idevice = { path = "../idevice", features = ["full"], default-features = false } tokio = { version = "1.43", features = ["full"] } diff --git a/tools/src/screenshot.rs b/tools/src/screenshot.rs new file mode 100644 index 0000000..bb6f617 --- /dev/null +++ b/tools/src/screenshot.rs @@ -0,0 +1,94 @@ +use clap::{Arg, Command}; +use idevice::{IdeviceService, RsdService, core_device_proxy::CoreDeviceProxy, rsd::RsdHandshake}; +use std::fs; + +mod common; + +#[tokio::main] +async fn main() { + env_logger::init(); + let matches = Command::new("screen_shot") + .about("take screenshot") + .arg( + Arg::new("host") + .long("host") + .value_name("HOST") + .help("IP address of the device"), + ) + .arg( + Arg::new("pairing_file") + .long("pairing-file") + .value_name("PATH") + .help("Path to the pairing file"), + ) + .arg( + Arg::new("udid") + .value_name("UDID") + .help("UDID of the device (overrides host/pairing file)") + .index(1), + ) + .arg( + Arg::new("output") + .short('o') + .long("output") + .value_name("FILE") + .help("Output file path for the screenshot (default: ./screenshot.png)") + .default_value("screenshot.png"), + ) + .arg( + Arg::new("about") + .long("about") + .help("Show about information") + .action(clap::ArgAction::SetTrue), + ) + .get_matches(); + + if matches.get_flag("about") { + print!("screen_shot - take screenshot from ios device"); + println!("Copyright (c) 2025 Jackson Coxson"); + return; + } + + let udid = matches.get_one::("udid"); + let host = matches.get_one::("host"); + let pairing_file = matches.get_one::("pairing_file"); + let output_path = matches.get_one::("output").unwrap(); + + let provider = + match common::get_provider(udid, host, pairing_file, "take_screenshot-jkcoxson").await { + Ok(p) => p, + Err(e) => { + eprintln!("{e}"); + return; + } + }; + let proxy = CoreDeviceProxy::connect(&*provider) + .await + .expect("no core proxy"); + let rsd_port = proxy.handshake.server_rsd_port; + + let adapter = proxy.create_software_tunnel().expect("no software tunnel"); + let mut adapter = adapter.to_async_handle(); + let stream = adapter.connect(rsd_port).await.expect("no RSD connect"); + + // Make the connection to RemoteXPC + let mut handshake = RsdHandshake::new(stream).await.unwrap(); + let mut ts_client = + idevice::dvt::remote_server::RemoteServerClient::connect_rsd(&mut adapter, &mut handshake) + .await + .expect("Failed to connect"); + ts_client.read_message(0).await.expect("no read??"); + + let mut ts_client = idevice::dvt::screenshot::ScreenshotClient::new(&mut ts_client) + .await + .expect("Unable to get channel for take screenshot"); + let res = ts_client + .take_screenshot() + .await + .expect("Failed to take screenshot"); + + match fs::write(output_path, &res) { + Ok(_) => println!("Screenshot saved to: {}", output_path), + Err(e) => eprintln!("Failed to write screenshot to file: {}", e), + } +}