change all PlistSocket to PlistSocketProxy to make socket release aut…

…omatically

.gitignore

@@ -144,3 +144,4 @@ ChangeLog
 usbmuxd-dumpdata
 .DS_Store
 tmp/
+*.pem

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

tidevice/_crash.py

@@ -4,15 +4,15 @@
 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)
 
 # 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

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)
 

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,

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):

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()

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)

tidevice/_safe_socket.py

@@ -2,7 +2,7 @@
 # codeskyblue 2020/06/03
 #
 
-__all__ = ['SafeStreamSocket', 'PlistSocket']
+__all__ = ['SafeStreamSocket', 'PlistSocket', 'PlistSocketProxy']
 
 import logging
 import os
@@ -202,31 +202,54 @@ 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()
 
     @property
     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()
 

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
 

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