cli,remote,service_connection: Implement userspace wrapper connect

Author: strobecat

pymobiledevice3/cli/cli_common.py

@@ -149,6 +149,7 @@ def wrap_callback_calling(**kwargs: dict) -> None:
         lockdown_service_provider = kwargs.pop('lockdown_service_provider', None)
         rsd_service_provider_manually = kwargs.pop('rsd_service_provider_manually', None)
         rsd_service_provider_using_tunneld = kwargs.pop('rsd_service_provider_using_tunneld', None)
+        del kwargs['rsd_use_userspace_wrapper']
         if lockdown_service_provider is not None:
             service_provider = lockdown_service_provider
         if rsd_service_provider_manually is not None:
@@ -224,23 +225,36 @@ def udid(self, ctx, param: str, value: str) -> Optional[LockdownClient]:
 class RSDCommand(BaseServiceProviderCommand):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.userspace_address = None
         self.params[:0] = [
+            RSDOption(('rsd_use_userspace_wrapper', '--userspace-wrapper'), type=(str, int),
+                      callback=self.userspace,
+                      mutually_exclusive=['rsd_service_provider_using_tunneld'],
+                      help='\b\nhostname and port number of userspace tunnel wrapper '
+                           '(as provided by a `start-tunnel` subcommand).'),
             RSDOption(('rsd_service_provider_manually', '--rsd'), type=(str, int), callback=self.rsd,
                       mutually_exclusive=['rsd_service_provider_using_tunneld'],
                       help='\b\n'
                            'RSD hostname and port number (as provided by a `start-tunnel` subcommand).'),
             RSDOption(('rsd_service_provider_using_tunneld', '--tunnel'), callback=self.tunneld,
-                      mutually_exclusive=['rsd_service_provider_manually'], envvar=TUNNEL_ENV_VAR,
+                      mutually_exclusive=[
+                          'rsd_service_provider_manually',
+                          'rsd_use_userspace_wrapper'
+                      ],
+                      envvar=TUNNEL_ENV_VAR,
                       help='\b\n'
                            'Either an empty string to force tunneld device selection, or a UDID of a tunneld '
                            'discovered device.\n'
                            'The string may be suffixed with :PORT in case tunneld is not serving at the default port.\n'
                            f'This option may also be transferred as an environment variable: {TUNNEL_ENV_VAR}')
         ]
 
+    def userspace(self, ctx, param: str, value: Optional[tuple[str, int]]) -> None:
+        self.userspace_address = value
+
     def rsd(self, ctx, param: str, value: Optional[tuple[str, int]]) -> Optional[RemoteServiceDiscoveryService]:
         if value is not None:
-            rsd = RemoteServiceDiscoveryService(value)
+            rsd = RemoteServiceDiscoveryService(value, userspace_address=self.userspace_address)
             asyncio.run(rsd.connect(), debug=True)
             self.service_provider = rsd
             return self.service_provider

pymobiledevice3/remote/remote_service_discovery.py

@@ -28,10 +28,16 @@ class RSDDevice:
 
 
 class RemoteServiceDiscoveryService(LockdownServiceProvider):
-    def __init__(self, address: tuple[str, int], name: Optional[str] = None) -> None:
+    def __init__(
+        self,
+        address: tuple[str, int],
+        userspace_address: "tuple[str, int] | None" = None,
+        name: Optional[str] = None
+    ) -> None:
         super().__init__()
         self.name = name
-        self.service = RemoteXPCConnection(address)
+        self.service = RemoteXPCConnection(address, userspace_address)
+        self.userspace_address = userspace_address
         self.peer_info: Optional[dict] = None
         self.lockdown: Optional[LockdownClient] = None
         self.all_values: Optional[dict] = None
@@ -75,7 +81,11 @@ def get_value(self, domain: Optional[str] = None, key: Optional[str] = None) ->
         return self.lockdown.get_value(domain, key)
 
     def start_lockdown_service_without_checkin(self, name: str) -> ServiceConnection:
-        return ServiceConnection.create_using_tcp(self.service.address[0], self.get_service_port(name))
+        return ServiceConnection.create_using_tcp(
+            self.service.address[0],
+            self.get_service_port(name),
+            userspace_address=self.userspace_address
+        )
 
     def start_lockdown_service(self, name: str, include_escrow_bag: bool = False) -> ServiceConnection:
         service = self.start_lockdown_service_without_checkin(name)
@@ -112,7 +122,10 @@ def start_lockdown_developer_service(self, name, include_escrow_bag: bool = Fals
             raise
 
     def start_remote_service(self, name: str) -> RemoteXPCConnection:
-        service = RemoteXPCConnection((self.service.address[0], self.get_service_port(name)))
+        service = RemoteXPCConnection(
+            (self.service.address[0], self.get_service_port(name)),
+            self.userspace_address
+        )
         return service
 
     def start_service(self, name: str) -> Union[RemoteXPCConnection, ServiceConnection]:

pymobiledevice3/remote/remotexpc.py

@@ -1,4 +1,5 @@
 import asyncio
+import socket
 import sys
 from asyncio import IncompleteReadError
 from typing import AsyncIterable, Optional
@@ -37,9 +38,10 @@
 
 
 class RemoteXPCConnection:
-    def __init__(self, address: tuple[str, int]):
+    def __init__(self, address: tuple[str, int], userspace_address: "tuple[str, int] | None" = None):
         self._previous_frame_data = b''
         self.address = address
+        self.userspace_address = userspace_address
         self.next_message_id: dict[int, int] = {ROOT_CHANNEL: 0, REPLY_CHANNEL: 0}
         self.peer_info = None
         self._reader: Optional[asyncio.StreamReader] = None
@@ -53,7 +55,15 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
         await self.close()
 
     async def connect(self) -> None:
-        self._reader, self._writer = await asyncio.open_connection(self.address[0], self.address[1])
+        if self.userspace_address:
+            self._reader, self._writer = await asyncio.open_connection(*self.userspace_address)
+            self._writer.write(
+                socket.inet_pton(socket.AF_INET6, self.address[0])
+                + self.address[1].to_bytes(4, "little")
+            )
+            await self._writer.drain()
+        else:
+            self._reader, self._writer = await asyncio.open_connection(self.address[0], self.address[1])
         try:
             await self._do_handshake()
         except Exception:  # noqa: E722

pymobiledevice3/service_connection.py

@@ -116,6 +116,7 @@ def __init__(self, sock: socket.socket, mux_device: MuxDevice = None):
 
     @staticmethod
     def create_using_tcp(hostname: str, port: int, keep_alive: bool = True,
+                         userspace_address: "None | tuple[str, int]" = None,
                          create_connection_timeout: int = DEFAULT_TIMEOUT) -> 'ServiceConnection':
         """
         Create a ServiceConnection using a TCP connection.
@@ -126,7 +127,11 @@ def create_using_tcp(hostname: str, port: int, keep_alive: bool = True,
         :param create_connection_timeout: The timeout for creating the connection.
         :return: A ServiceConnection object.
         """
-        sock = socket.create_connection((hostname, port), timeout=create_connection_timeout)
+        if userspace_address:
+            sock = socket.create_connection(userspace_address, timeout=create_connection_timeout)
+            sock.sendall(socket.inet_pton(socket.AF_INET6, hostname) + port.to_bytes(4, "little"))
+        else:
+            sock = socket.create_connection((hostname, port), timeout=create_connection_timeout)
         sock.settimeout(None)
         if keep_alive:
             OSUTIL.set_keepalive(sock)