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 <jkcoxson@gmail.com>

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<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Self, IdeviceError> {
-        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<u8> = 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 {

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<Box<dyn ReadWrite>> {
     fn rsd_service_name() -> std::borrow::Cow<'static, str> {

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<R>) -> Result<Self, IdeviceError> {
+        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<u8>)` - 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<Vec<u8>, 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),
+        }
+    }
+}

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"] }

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::<String>("udid");
+    let host = matches.get_one::<String>("host");
+    let pairing_file = matches.get_one::<String>("pairing_file");
+    let output_path = matches.get_one::<String>("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),
+    }
+}