From 32b8de662421bcd4ca9727a9cd0387505084e302 Mon Sep 17 00:00:00 2001 From: codeskyblue Date: Fri, 17 Jun 2022 23:50:09 +0800 Subject: [PATCH] change all PlistSocket to PlistSocketProxy to make socket release automatically --- .gitignore | 1 + tidevice/__main__.py | 2 +- tidevice/_crash.py | 4 +-- tidevice/_device.py | 69 +++++++++++++++++++++------------------ tidevice/_imagemounter.py | 8 ++--- tidevice/_installation.py | 4 +-- tidevice/_instruments.py | 4 +-- tidevice/_relay.py | 4 +-- tidevice/_safe_socket.py | 39 +++++++++++++++++----- tidevice/_sync.py | 4 +-- tidevice/_usbmux.py | 19 ++++++----- 11 files changed, 95 insertions(+), 63 deletions(-) diff --git a/.gitignore b/.gitignore index edd96a4..6a07200 100644 --- a/.gitignore +++ b/.gitignore @@ -144,3 +144,4 @@ ChangeLog usbmuxd-dumpdata .DS_Store tmp/ +*.pem diff --git a/tidevice/__main__.py b/tidevice/__main__.py index 9b42438..6ad2d8c 100644 --- a/tidevice/__main__.py +++ b/tidevice/__main__.py @@ -422,7 +422,7 @@ def cmd_syslog(args: argparse.Namespace): s = d.start_service("com.apple.syslog_relay") try: while True: - text = s.recv().decode('utf-8') + text = s.psock.recv().decode('utf-8') print(text, end='', flush=True) except (BrokenPipeError, IOError): # Python flushes standard streams on exit; redirect remaining output diff --git a/tidevice/_crash.py b/tidevice/_crash.py index 9ee3386..41847f8 100644 --- a/tidevice/_crash.py +++ b/tidevice/_crash.py @@ -4,7 +4,7 @@ from ._sync import Sync import logging from ._proto import LOG -from ._safe_socket import PlistSocket +from ._safe_socket import PlistSocketProxy logger = logging.getLogger(LOG.main) @@ -12,7 +12,7 @@ # Ref: https://github.com/libimobiledevice/libimobiledevice/blob/master/tools/idevicecrashreport.c class CrashManager: - def __init__(self, copy_conn: PlistSocket): + def __init__(self, copy_conn: PlistSocketProxy): self._afc = Sync(copy_conn) @property diff --git a/tidevice/_device.py b/tidevice/_device.py index 8d0cb5b..52b0296 100644 --- a/tidevice/_device.py +++ b/tidevice/_device.py @@ -283,7 +283,7 @@ def create_inner_connection( self, port: int = LOCKDOWN_PORT, # 0xf27e, _ssl: bool = False, - ssl_dial_only: bool = False) -> PlistSocket: + ssl_dial_only: bool = False) -> PlistSocketProxy: _port = socket.htons(port) # Same as: ((port & 0xff) << 8) | (port >> 8) del (port) @@ -303,17 +303,21 @@ def create_inner_connection( self._usbmux._check(data) logger.debug("connected to port: %d", _port) if _ssl: - conn.switch_to_ssl(self.ssl_pemfile_path) + psock = conn.psock + psock.switch_to_ssl(self.ssl_pemfile_path) if ssl_dial_only: - conn.ssl_unwrap() + psock.ssl_unwrap() return conn @contextlib.contextmanager - def create_session(self) -> PlistSocket: + def create_session(self) -> PlistSocketProxy: """ Create session inside SSLContext """ - with self.create_inner_connection() as s: # 62078=0xf27e + with self.create_inner_connection() as _s: # 62078=0xf27e + s: PlistSocketProxy = _s + del(_s) + data = s.send_recv_packet({"Request": "QueryType"}) # Expect: {'Request': 'QueryType', 'Type': 'com.apple.mobile.lockdown'} assert data['Type'] == LockdownService.MobileLockdown @@ -351,7 +355,7 @@ def create_session(self) -> PlistSocket: if data['EnableSessionSSL']: # tempfile.NamedTemporaryFile is not working well on windows # See: https://stackoverflow.com/questions/6416782/what-is-namedtemporaryfile-useful-for-on-windows - s.switch_to_ssl(self.ssl_pemfile_path) + s.psock.switch_to_ssl(self.ssl_pemfile_path) yield s @@ -481,13 +485,13 @@ def get_crashmanager(self) -> CrashManager: # read "ping" message which indicates the crash logs have been moved to a safe harbor move_conn = self.start_service(LockdownService.CRASH_REPORT_MOVER_SERVICE) ack = b'ping\x00' - if ack != move_conn.recvall(len(ack)): + if ack != move_conn.psock.recvall(len(ack)): raise ServiceError("ERROR: Crash logs could not be moved. Connection interrupted") copy_conn = self.start_service(LockdownService.CRASH_REPORT_COPY_MOBILE_SERVICE) return CrashManager(copy_conn) - def start_service(self, name: str) -> PlistSocket: + def start_service(self, name: str) -> PlistSocketProxy: try: return self._unsafe_start_service(name) except MuxServiceError: @@ -496,8 +500,11 @@ def start_service(self, name: str) -> PlistSocket: time.sleep(.5) return self._unsafe_start_service(name) - def _unsafe_start_service(self, name: str) -> PlistSocket: - with self.create_session() as s: + def _unsafe_start_service(self, name: str) -> PlistSocketProxy: + with self.create_session() as _s: + s: PlistSocketProxy = _s + del(_s) + data = s.send_recv_packet({ "Request": "StartService", "Service": name, @@ -535,31 +542,31 @@ def screenshot(self) -> Image.Image: def iter_screenshot(self) -> Iterator[Image.Image]: """ take screenshot infinite """ - with self.start_service(LockdownService.MobileScreenshotr) as conn: - version_exchange = conn.recv_packet() - # Expect recv: ['DLMessageVersionExchange', 300, 0] + conn = self.start_service(LockdownService.MobileScreenshotr) + version_exchange = conn.recv_packet() + # Expect recv: ['DLMessageVersionExchange', 300, 0] + data = conn.send_recv_packet([ + 'DLMessageVersionExchange', 'DLVersionsOk', version_exchange[1] + ]) + # Expect recv: ['DLMessageDeviceReady'] + assert data[0] == 'DLMessageDeviceReady' + + while True: + # code will be blocked here until next(..) called data = conn.send_recv_packet([ - 'DLMessageVersionExchange', 'DLVersionsOk', version_exchange[1] + 'DLMessageProcessMessage', { + 'MessageType': 'ScreenShotRequest' + } ]) - # Expect recv: ['DLMessageDeviceReady'] - assert data[0] == 'DLMessageDeviceReady' - - while True: - # code will be blocked here until next(..) called - data = conn.send_recv_packet([ - 'DLMessageProcessMessage', { - 'MessageType': 'ScreenShotRequest' - } - ]) - # Expect recv: ['DLMessageProcessMessage', {'MessageType': 'ScreenShotReply', ScreenShotData': b'\x89PNG\r\n\x...'}] - assert len(data) == 2 and data[0] == 'DLMessageProcessMessage' - assert isinstance(data[1], dict) - assert data[1]['MessageType'] == "ScreenShotReply" + # Expect recv: ['DLMessageProcessMessage', {'MessageType': 'ScreenShotReply', ScreenShotData': b'\x89PNG\r\n\x...'}] + assert len(data) == 2 and data[0] == 'DLMessageProcessMessage' + assert isinstance(data[1], dict) + assert data[1]['MessageType'] == "ScreenShotReply" - png_data = data[1]['ScreenShotData'] + png_data = data[1]['ScreenShotData'] - yield pil_imread(png_data) + yield pil_imread(png_data) @property def name(self): @@ -813,7 +820,7 @@ def _launch_app_runner(self, env: dict = {}, target_app_bundle_id: str = None, logger: logging.Logger = logging, - quit_event: threading.Event = None) -> typing.Tuple[PlistSocket, int]: # pid + quit_event: threading.Event = None) -> typing.Tuple[ServiceInstruments, int]: # pid logger = logging.getLogger(LOG.xctest) diff --git a/tidevice/_imagemounter.py b/tidevice/_imagemounter.py index b7e3180..8764ff5 100644 --- a/tidevice/_imagemounter.py +++ b/tidevice/_imagemounter.py @@ -11,7 +11,7 @@ import retry import requests -from ._safe_socket import PlistSocketProperty +from ._safe_socket import PlistSocketProxy from ._utils import get_app_dir, logger from .exceptions import MuxError, MuxServiceError @@ -88,7 +88,7 @@ def cache_developer_image(version: str) -> str: return image_zip_path -class ImageMounter(PlistSocketProperty): +class ImageMounter(PlistSocketProxy): SERVICE_NAME = "com.apple.mobile.mobile_image_mounter" def prepare(self): @@ -101,7 +101,7 @@ def lookup(self, image_type="Developer") -> List[bytes]: """ Check image signature """ - ret = self.psock.send_recv_packet({ + ret = self.send_recv_packet({ "Command": "LookupImage", "ImageType": image_type, }) @@ -144,7 +144,7 @@ def mount_fileobj(self, signature_content: bytes, image_type: str = "Developer"): - ret = self.psock.send_recv_packet({ + ret = self.send_recv_packet({ "Command": "ReceiveBytes", "ImageSignature": signature_content, "ImageSize": image_size, diff --git a/tidevice/_installation.py b/tidevice/_installation.py index c8b3aea..185c88e 100644 --- a/tidevice/_installation.py +++ b/tidevice/_installation.py @@ -22,13 +22,13 @@ from typing import Optional from ._proto import LOG -from ._safe_socket import PlistSocketProperty +from ._safe_socket import PlistSocketProxy from .exceptions import ServiceError logger = logging.getLogger(LOG.main) -class Installation(PlistSocketProperty): +class Installation(PlistSocketProxy): SERVICE_NAME = "com.apple.mobile.installation_proxy" def install(self, bundle_id: str, target_path: str): diff --git a/tidevice/_instruments.py b/tidevice/_instruments.py index a5aa168..2276ade 100644 --- a/tidevice/_instruments.py +++ b/tidevice/_instruments.py @@ -29,7 +29,7 @@ from . import bplist from . import struct2 as ct from ._proto import LOG, InstrumentsService -from ._safe_socket import PlistSocketProperty +from ._safe_socket import PlistSocketProxy from .exceptions import MuxError, ServiceError logger = logging.getLogger(LOG.xctest) @@ -267,7 +267,7 @@ def append_obj(self, obj): # else: # self.append_obj(v) -class DTXService(PlistSocketProperty): +class DTXService(PlistSocketProxy): def prepare(self): super().prepare() diff --git a/tidevice/_relay.py b/tidevice/_relay.py index d83739a..597753c 100644 --- a/tidevice/_relay.py +++ b/tidevice/_relay.py @@ -15,7 +15,7 @@ from ._device import Device from ._hexdump import hexdump -from ._safe_socket import PlistSocket +from ._safe_socket import PlistSocketProxy from .exceptions import MuxReplyError @@ -51,7 +51,7 @@ def _pipe_twoway(self, _in: IOStream, out: IOStream, plconn): io_loop.add_callback(self._pipe_stream, _in, out, plconn) io_loop.add_callback(self._pipe_stream, out, _in, plconn) - async def _pipe_stream(self, _in: IOStream, out: IOStream, plconn: PlistSocket): + async def _pipe_stream(self, _in: IOStream, out: IOStream, plconn: PlistSocketProxy): while not _in.closed(): try: data = await _in.read_bytes(10240, partial=True) diff --git a/tidevice/_safe_socket.py b/tidevice/_safe_socket.py index 0f0d536..3321b1d 100644 --- a/tidevice/_safe_socket.py +++ b/tidevice/_safe_socket.py @@ -2,7 +2,7 @@ # codeskyblue 2020/06/03 # -__all__ = ['SafeStreamSocket', 'PlistSocket'] +__all__ = ['SafeStreamSocket', 'PlistSocket', 'PlistSocketProxy'] import logging import os @@ -202,15 +202,16 @@ def recv_packet(self, header_size=None) -> dict: logger.debug("RECV(%d): %s", self.id, payload) return payload - def send_recv_packet(self, payload: dict, timeout: float = 10.0) -> dict: - with set_socket_timeout(self.get_socket(), timeout): - self.send_packet(payload) - return self.recv_packet() +class PlistSocketProxy: + def __init__(self, psock: typing.Union[PlistSocket, "PlistSocketProxy"]): + if isinstance(psock, PlistSocketProxy): + self._psock = psock._psock + psock._finalizer.detach() + else: + assert isinstance(psock, PlistSocket) + self._psock = psock -class PlistSocketProperty: - def __init__(self, psock: PlistSocket): - self._psock = psock self._finalizer = weakref.finalize(self, self._psock.close) self.prepare() @@ -218,15 +219,37 @@ def __init__(self, psock: PlistSocket): def psock(self) -> PlistSocket: return self._psock + @property + def name(self) -> str: + return self.psock.name + + @name.setter + def name(self, new_name: str): + self.psock.name = new_name + def prepare(self): pass + def get_socket(self) -> socket.socket: + return self.psock.get_socket() + def send_packet(self, payload: dict, message_type: int = 8): return self.psock.send_packet(payload, message_type) def recv_packet(self, header_size=None) -> dict: return self.psock.recv_packet(header_size) + def send_recv_packet(self, payload: dict, timeout: float = 10.0) -> dict: + with set_socket_timeout(self.psock.get_socket(), timeout): + self.send_packet(payload) + return self.recv_packet() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + def close(self): self._finalizer() diff --git a/tidevice/_sync.py b/tidevice/_sync.py index 7161ee3..bdf7d19 100644 --- a/tidevice/_sync.py +++ b/tidevice/_sync.py @@ -16,7 +16,7 @@ from . import bplist from . import struct2 as ct from ._proto import * -from ._safe_socket import PlistSocketProperty +from ._safe_socket import PlistSocketProxy from ._utils import pathjoin from .exceptions import MuxError, MuxServiceError @@ -38,7 +38,7 @@ logger = logging.getLogger(PROGRAM_NAME) -class Sync(PlistSocketProperty): +class Sync(PlistSocketProxy): def prepare(self): self.__tag = -1 diff --git a/tidevice/_usbmux.py b/tidevice/_usbmux.py index 614618e..3f215de 100644 --- a/tidevice/_usbmux.py +++ b/tidevice/_usbmux.py @@ -12,7 +12,7 @@ from ._types import DeviceInfo, ConnectionType from ._proto import PROGRAM_NAME, UsbmuxReplyCode -from ._safe_socket import PlistSocket +from ._safe_socket import PlistSocket, PlistSocketProxy from .exceptions import * # pragma warning disables S2208 @@ -40,16 +40,17 @@ def _next_tag(self) -> int: self.__tag += 1 return self.__tag - def create_connection(self) -> PlistSocket: - return PlistSocket(self.__address, self._next_tag()) + def create_connection(self) -> PlistSocketProxy: + psock = PlistSocket(self.__address, self._next_tag()) + return PlistSocketProxy(psock) def send_recv(self, payload: dict, timeout: float = None) -> dict: - with self.create_connection() as s: - s.get_socket().settimeout(timeout) - s.send_packet(payload) - recv_data = s.recv_packet() - self._check(recv_data) - return recv_data + s = self.create_connection() + s.get_socket().settimeout(timeout) + s.send_packet(payload) + recv_data = s.recv_packet() + self._check(recv_data) + return recv_data def device_list(self) -> typing.List[DeviceInfo]: """