lockdown: add retry_create_using_usbmux()

pymobiledevice3/lockdown.py

@@ -15,6 +15,7 @@
 from ssl import SSLError, SSLZeroReturnError
 from typing import Optional
 
+import construct
 from cryptography import x509
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
@@ -26,10 +27,11 @@
 from pymobiledevice3.bonjour import DEFAULT_BONJOUR_TIMEOUT, browse_mobdev2
 from pymobiledevice3.ca import ca_do_everything
 from pymobiledevice3.common import get_home_folder
-from pymobiledevice3.exceptions import CannotStopSessionError, ConnectionTerminatedError, FatalPairingError, \
-    GetProhibitedError, IncorrectModeError, InvalidConnectionError, InvalidHostIDError, InvalidServiceError, \
-    LockdownError, MissingValueError, NotPairedError, PairingDialogResponsePendingError, PairingError, \
-    PasswordRequiredError, SetProhibitedError, StartServiceError, UserDeniedPairingError
+from pymobiledevice3.exceptions import BadDevError, CannotStopSessionError, ConnectionFailedError, \
+    ConnectionTerminatedError, DeviceNotFoundError, FatalPairingError, GetProhibitedError, IncorrectModeError, \
+    InvalidConnectionError, InvalidHostIDError, InvalidServiceError, LockdownError, MissingValueError, \
+    NoDeviceConnectedError, NotPairedError, PairingDialogResponsePendingError, PairingError, PasswordRequiredError, \
+    SetProhibitedError, StartServiceError, UserDeniedPairingError
 from pymobiledevice3.irecv_devices import IRECV_DEVICES
 from pymobiledevice3.lockdown_service_provider import LockdownServiceProvider
 from pymobiledevice3.pair_records import create_pairing_records_cache_folder, generate_host_id, \
@@ -747,6 +749,23 @@ def create_using_usbmux(serial: str = None, identifier: str = None, label: str =
         autopair=autopair, usbmux_address=usbmux_address)
 
 
+def retry_create_using_usbmux(retry_timeout: Optional[float] = None, **kwargs) -> UsbmuxLockdownClient:
+    """
+    Repeatedly retry to create a UsbmuxLockdownClient instance while dismissing different errors that might occur
+    while device is rebooting
+
+    :param retry_timeout: Retry timeout in seconds or None for no timeout
+    :return: UsbmuxLockdownClient instance
+    """
+    start = time.time()
+    while (retry_timeout is None) or (time.time() - start < retry_timeout):
+        try:
+            return create_using_usbmux(**kwargs)
+        except (NoDeviceConnectedError, ConnectionFailedError, BadDevError, OSError, construct.core.StreamError,
+                DeviceNotFoundError):
+            pass
+
+
 def create_using_tcp(hostname: str, identifier: str = None, label: str = DEFAULT_LABEL, autopair: bool = True,
                      pair_timeout: float = None, local_hostname: str = None, pair_record: Optional[dict] = None,
                      pairing_records_cache_folder: Path = None, port: int = SERVICE_PORT,

pymobiledevice3/services/amfi.py

@@ -1,11 +1,9 @@
 #!/usr/bin/env python3
 import logging
 
-import construct
-
-from pymobiledevice3.exceptions import AmfiError, BadDevError, ConnectionFailedError, DeveloperModeError, \
-    DeviceHasPasscodeSetError, DeviceNotFoundError, NoDeviceConnectedError, PyMobileDevice3Exception
-from pymobiledevice3.lockdown import LockdownClient, create_using_usbmux
+from pymobiledevice3.exceptions import AmfiError, DeveloperModeError, DeviceHasPasscodeSetError, \
+    PyMobileDevice3Exception
+from pymobiledevice3.lockdown import LockdownClient, retry_create_using_usbmux
 from pymobiledevice3.services.heartbeat import HeartbeatService
 
 
@@ -54,14 +52,7 @@ def enable_developer_mode(self, enable_post_restart=True):
         except ConnectionAbortedError:
             self._logger.debug('device disconnected, awaiting reconnect')
 
-        while True:
-            try:
-                self._lockdown = create_using_usbmux(self._lockdown.udid)
-                break
-            except (NoDeviceConnectedError, ConnectionFailedError, BadDevError, OSError, construct.core.StreamError,
-                    DeviceNotFoundError):
-                pass
-
+        self._lockdown = retry_create_using_usbmux(None, serial=self._lockdown.udid)
         self.enable_developer_mode_post_restart()
 
     def enable_developer_mode_post_restart(self):