pymobiledevice3 4.14.6__py3-none-any.whl → 7.0.6__py3-none-any.whl

pymobiledevice3/remote/remotexpc.py

@@ -1,20 +1,34 @@
 import asyncio
+import contextlib
 import sys
 from asyncio import IncompleteReadError
-from collections.abc import Generator
+from collections.abc import AsyncIterable
 from typing import Optional
 
 import IPython
 import nest_asyncio
 from construct import StreamError
-from hyperframe.frame import DataFrame, Frame, GoAwayFrame, HeadersFrame, RstStreamFrame, SettingsFrame, \
-    WindowUpdateFrame
+from hyperframe.frame import (
+    DataFrame,
+    Frame,
+    GoAwayFrame,
+    HeadersFrame,
+    RstStreamFrame,
+    SettingsFrame,
+    WindowUpdateFrame,
+)
 from pygments import formatters, highlight, lexers
 from traitlets.config import Config
 
-from pymobiledevice3.exceptions import StreamClosedError
-from pymobiledevice3.remote.xpc_message import XpcFlags, XpcInt64Type, XpcUInt64Type, XpcWrapper, create_xpc_wrapper, \
-    decode_xpc_object
+from pymobiledevice3.exceptions import ProtocolError, StreamClosedError
+from pymobiledevice3.remote.xpc_message import (
+    XpcFlags,
+    XpcInt64Type,
+    XpcUInt64Type,
+    XpcWrapper,
+    create_xpc_wrapper,
+    decode_xpc_object,
+)
 
 # Extracted by sniffing `remoted` traffic via Wireshark
 DEFAULT_SETTINGS_MAX_CONCURRENT_STREAMS = 100
@@ -22,11 +36,13 @@ DEFAULT_SETTINGS_INITIAL_WINDOW_SIZE = 1048576
 DEFAULT_WIN_SIZE_INCR = 983041
 
 FRAME_HEADER_SIZE = 9
-HTTP2_MAGIC = b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n'
+HTTP2_MAGIC = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
 
 ROOT_CHANNEL = 1
 REPLY_CHANNEL = 3
 
+FIRST_REPLY_TIMEOUT = 3
+
 SHELL_USAGE = """
 # This shell allows you to communicate directly with every RemoteXPC service.
 
@@ -37,14 +53,14 @@ resp = await client.send_receive_request({"Command": "DoSomething"})
 
 class RemoteXPCConnection:
     def __init__(self, address: tuple[str, int]):
-        self._previous_frame_data = b''
+        self._previous_frame_data = b""
         self.address = address
         self.next_message_id: dict[int, int] = {ROOT_CHANNEL: 0, REPLY_CHANNEL: 0}
         self.peer_info = None
         self._reader: Optional[asyncio.StreamReader] = None
         self._writer: Optional[asyncio.StreamWriter] = None
 
-    async def __aenter__(self) -> 'RemoteXPCConnection':
+    async def __aenter__(self) -> "RemoteXPCConnection":
         await self.connect()
         return self
 
@@ -53,33 +69,36 @@ class RemoteXPCConnection:
 
     async def connect(self) -> None:
         self._reader, self._writer = await asyncio.open_connection(self.address[0], self.address[1])
-        await self._do_handshake()
+        try:
+            await self._do_handshake()
+        except Exception:
+            await self.close()
+            raise
 
     async def close(self) -> None:
         if self._writer is None:
             return
         self._writer.close()
-        try:
+        with contextlib.suppress(ConnectionResetError):
             await self._writer.wait_closed()
-        except ConnectionResetError:
-            pass
         self._writer = None
         self._reader = None
 
     async def send_request(self, data: dict, wanting_reply: bool = False) -> None:
         xpc_wrapper = create_xpc_wrapper(
-            data, message_id=self.next_message_id[ROOT_CHANNEL], wanting_reply=wanting_reply)
+            data, message_id=self.next_message_id[ROOT_CHANNEL], wanting_reply=wanting_reply
+        )
         self._writer.write(DataFrame(stream_id=ROOT_CHANNEL, data=xpc_wrapper).serialize())
         await self._writer.drain()
 
-    async def iter_file_chunks(self, total_size: int, file_idx: int = 0) -> Generator[bytes, None, None]:
+    async def iter_file_chunks(self, total_size: int, file_idx: int = 0) -> AsyncIterable[bytes]:
         stream_id = (file_idx + 1) * 2
         await self._open_channel(stream_id, XpcFlags.FILE_TX_STREAM_RESPONSE)
         size = 0
         while size < total_size:
             frame = await self._receive_next_data_frame()
 
-            if 'END_STREAM' in frame.flags:
+            if "END_STREAM" in frame.flags:
                 continue
 
             if frame.stream_id != stream_id:
@@ -87,12 +106,12 @@ class RemoteXPCConnection:
                 if xpc_wrapper.flags.FILE_TX_STREAM_REQUEST:
                     continue
 
-            assert frame.stream_id == stream_id, f'got {frame.stream_id} instead of {stream_id}'
+            assert frame.stream_id == stream_id, f"got {frame.stream_id} instead of {stream_id}"
             size += len(frame.data)
             yield frame.data
 
     async def receive_file(self, total_size: int) -> bytes:
-        buf = b''
+        buf = b""
         async for chunk in self.iter_file_chunks(total_size):
             buf += chunk
         return buf
@@ -102,7 +121,7 @@ class RemoteXPCConnection:
             frame = await self._receive_next_data_frame()
             try:
                 xpc_message = XpcWrapper.parse(self._previous_frame_data + frame.data).message
-                self._previous_frame_data = b''
+                self._previous_frame_data = b""
             except StreamError:
                 self._previous_frame_data += frame.data
                 continue
@@ -119,46 +138,56 @@ class RemoteXPCConnection:
 
     def shell(self) -> None:
         nest_asyncio.apply(asyncio.get_running_loop())
-        sys.argv = ['a']
+        sys.argv = ["a"]
         config = Config()
-        config.InteractiveShellApp.exec_lines = ['%autoawait asyncio']
-        print(highlight(SHELL_USAGE, lexers.PythonLexer(),
-                        formatters.Terminal256Formatter(style='native')))
-        IPython.start_ipython(config=config, user_ns={
-            'client': self,
-            'XpcInt64Type': XpcInt64Type,
-            'XpcUInt64Type': XpcUInt64Type,
-        })
+        config.InteractiveShellApp.exec_lines = ["%autoawait asyncio"]
+        print(highlight(SHELL_USAGE, lexers.PythonLexer(), formatters.Terminal256Formatter(style="native")))
+        IPython.start_ipython(
+            config=config,
+            user_ns={
+                "client": self,
+                "XpcInt64Type": XpcInt64Type,
+                "XpcUInt64Type": XpcUInt64Type,
+            },
+        )
 
     async def _do_handshake(self) -> None:
         self._writer.write(HTTP2_MAGIC)
         await self._writer.drain()
 
         # send h2 headers
-        await self._send_frame(SettingsFrame(settings={
-            SettingsFrame.MAX_CONCURRENT_STREAMS: DEFAULT_SETTINGS_MAX_CONCURRENT_STREAMS,
-            SettingsFrame.INITIAL_WINDOW_SIZE: DEFAULT_SETTINGS_INITIAL_WINDOW_SIZE,
-        }))
+        await self._send_frame(
+            SettingsFrame(
+                settings={
+                    SettingsFrame.MAX_CONCURRENT_STREAMS: DEFAULT_SETTINGS_MAX_CONCURRENT_STREAMS,
+                    SettingsFrame.INITIAL_WINDOW_SIZE: DEFAULT_SETTINGS_INITIAL_WINDOW_SIZE,
+                }
+            )
+        )
         await self._send_frame(WindowUpdateFrame(stream_id=0, window_increment=DEFAULT_WIN_SIZE_INCR))
-        await self._send_frame(HeadersFrame(stream_id=ROOT_CHANNEL, flags=['END_HEADERS']))
+        await self._send_frame(HeadersFrame(stream_id=ROOT_CHANNEL, flags=["END_HEADERS"]))
 
         # send first actual requests
         await self.send_request({})
-        await self._send_frame(DataFrame(stream_id=ROOT_CHANNEL,
-                                         data=XpcWrapper.build({'size': 0, 'flags': 0x0201, 'payload': None})))
+        await self._send_frame(
+            DataFrame(stream_id=ROOT_CHANNEL, data=XpcWrapper.build({"size": 0, "flags": 0x0201, "payload": None}))
+        )
         self.next_message_id[ROOT_CHANNEL] += 1
         await self._open_channel(REPLY_CHANNEL, XpcFlags.INIT_HANDSHAKE)
         self.next_message_id[REPLY_CHANNEL] += 1
 
-        assert isinstance(await self._receive_frame(), SettingsFrame)
+        settings_frame = await asyncio.wait_for(self._receive_frame(), FIRST_REPLY_TIMEOUT)
+        if not isinstance(settings_frame, SettingsFrame):
+            raise ProtocolError(f"Got unexpected frame: {settings_frame} instead of a SettingsFrame")
 
-        await self._send_frame(SettingsFrame(flags=['ACK']))
+        await self._send_frame(SettingsFrame(flags=["ACK"]))
 
     async def _open_channel(self, stream_id: int, flags: int) -> None:
         flags |= XpcFlags.ALWAYS_SET
-        await self._send_frame(HeadersFrame(stream_id=stream_id, flags=['END_HEADERS']))
+        await self._send_frame(HeadersFrame(stream_id=stream_id, flags=["END_HEADERS"]))
         await self._send_frame(
-            DataFrame(stream_id=stream_id, data=XpcWrapper.build({'size': 0, 'flags': flags, 'payload': None})))
+            DataFrame(stream_id=stream_id, data=XpcWrapper.build({"size": 0, "flags": flags, "payload": None}))
+        )
 
     async def _send_frame(self, frame: Frame) -> None:
         self._writer.write(frame.serialize())
@@ -169,9 +198,9 @@ class RemoteXPCConnection:
             frame = await self._receive_frame()
 
             if isinstance(frame, GoAwayFrame):
-                raise StreamClosedError(f'Got {frame}')
+                raise StreamClosedError(f"Got {frame}")
             if isinstance(frame, RstStreamFrame):
-                raise StreamClosedError(f'Got {frame}')
+                raise StreamClosedError(f"Got {frame}")
             if not isinstance(frame, DataFrame):
                 continue
 
@@ -188,11 +217,11 @@ class RemoteXPCConnection:
         return frame
 
     async def _recvall(self, size: int) -> bytes:
-        data = b''
+        data = b""
         while len(data) < size:
             try:
                 chunk = await self._reader.readexactly(size - len(data))
-            except IncompleteReadError:
-                raise ConnectionAbortedError()
+            except IncompleteReadError as e:
+                raise ConnectionAbortedError() from e
             data += chunk
         return data