Merge pull request #1334 from doronz88/feature/reconnect

Feature/reconnect

pymobiledevice3/__main__.py

@@ -18,6 +18,7 @@
     MessageNotSupportedError, MissingValueError, NoDeviceConnectedError, NotEnoughDiskSpaceError, NotPairedError, \
     OSNotSupportedError, PairingDialogResponsePendingError, PasswordRequiredError, QuicProtocolNotSupportedError, \
     RSDRequiredError, SetProhibitedError, TunneldConnectionError, UserDeniedPairingError
+from pymobiledevice3.lockdown import retry_create_using_usbmux
 from pymobiledevice3.osu.os_utils import get_os_utils
 
 coloredlogs.install(level=logging.INFO)
@@ -89,6 +90,9 @@
     'install-completions': 'completions',
 }
 
+# Set if used the `--reconnect` option
+RECONNECT = False
+
 
 class Pmd3Cli(click.Group):
     def list_commands(self, ctx):
@@ -163,14 +167,16 @@ def load_all_commands() -> list[str]:
 
 
 @click.command(cls=Pmd3Cli, context_settings=CONTEXT_SETTINGS)
-def cli():
+@click.option('--reconnect', is_flag=True, default=False, help='Reconnect to device when disconnected.')
+def cli(reconnect: bool) -> None:
     """
     \b
     Interact with a connected iDevice (iPhone, iPad, ...)
     For more information please look at:
         https://github.com/doronz88/pymobiledevice3
     """
-    pass
+    global RECONNECT
+    RECONNECT = reconnect
 
 
 def main() -> None:
@@ -180,6 +186,10 @@ def main() -> None:
         logger.error('Device is not connected')
     except ConnectionAbortedError:
         logger.error('Device was disconnected')
+        if RECONNECT:
+            lockdown = retry_create_using_usbmux(None)
+            lockdown.close()
+            cli()
     except NotPairedError:
         logger.error('Device is not paired')
     except UserDeniedPairingError:

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

pymobiledevice3/services/crash_reports.py

@@ -3,13 +3,15 @@
 import re
 import time
 from collections.abc import Generator
+from json import JSONDecodeError
 from typing import Callable, Optional
 
 from pycrashreport.crash_report import get_crash_report_from_buf
 from xonsh.built_ins import XSH
 from xonsh.cli_utils import Annotated, Arg
 
-from pymobiledevice3.exceptions import AfcException, NotificationTimeoutError, SysdiagnoseTimeoutError
+from pymobiledevice3.exceptions import AfcException, AfcFileNotFoundError, NotificationTimeoutError, \
+    SysdiagnoseTimeoutError
 from pymobiledevice3.lockdown import LockdownClient
 from pymobiledevice3.lockdown_service_provider import LockdownServiceProvider
 from pymobiledevice3.services.afc import AfcService, AfcShell, path_completer
@@ -124,8 +126,14 @@ def watch(self, name: str = None, raw: bool = False) -> Generator[str, None, Non
             if posixpath.splitext(filename)[-1] not in ('.ips', '.panic'):
                 continue
 
-            crash_report_raw = self.afc.get_file_contents(filename).decode()
-            crash_report = get_crash_report_from_buf(crash_report_raw, filename=filename)
+            while True:
+                try:
+                    crash_report_raw = self.afc.get_file_contents(filename).decode()
+                    crash_report = get_crash_report_from_buf(crash_report_raw, filename=filename)
+                    break
+                except (AfcFileNotFoundError, JSONDecodeError):
+                    # Sometimes we have to wait for the file to be readable
+                    pass
 
             if name is None or crash_report.name == name:
                 if raw: