Merge pull request #202 from doronz88/bugfix/multiple_udids

usbmux: bugfix: use select_device to choose device correctly

pymobiledevice3/cli/developer.py

@@ -739,16 +739,16 @@ def debugserver_applist(lockdown: LockdownClient):
 
 @debugserver.command('start-server', cls=Command)
 @click.argument('local_port', type=click.INT)
-def debugserver_shell(lockdown: LockdownClient, local_port):
+def debugserver_start_server(lockdown: LockdownClient, local_port):
     """
     start a debugserver at remote listening on a given port locally.
 
     Please note the connection must be done soon afterwards using your own lldb client.
     This can be done using the following commands within lldb shell:
 
-    - platform select remote-ios
+    (lldb) platform select remote-ios
 
-    - platform connect connect://localhost:<local_port>
+    (lldb) platform connect connect://localhost:<local_port>
     """
     attr = lockdown.get_service_connection_attributes('com.apple.debugserver.DVTSecureSocketProxy')
     TcpForwarder(lockdown, local_port, attr['Port'], attr.get('EnableServiceSSL', False)).start()

pymobiledevice3/cli/list_devices.py

@@ -24,7 +24,7 @@ def list_devices(color, usb):
             if ':' in udid:
                 continue
 
-        lockdown = LockdownClient(udid)
+        lockdown = LockdownClient(udid, autopair=False)
         connected_devices.append(lockdown.all_values)
 
     print_json(connected_devices, colored=color)

pymobiledevice3/cli/restore.py

@@ -7,10 +7,11 @@
 import click
 from pygments import highlight, lexers, formatters
 
+from pymobiledevice3 import usbmux
 from pymobiledevice3.cli.cli_common import print_json, set_verbosity
 from pymobiledevice3.exceptions import IncorrectModeError
 from pymobiledevice3.irecv import IRecv
-from pymobiledevice3.lockdown import list_devices, LockdownClient
+from pymobiledevice3.lockdown import LockdownClient
 from pymobiledevice3.restore.device import Device
 from pymobiledevice3.restore.recovery import Recovery
 from pymobiledevice3.restore.restore import Restore
@@ -40,9 +41,9 @@ def device(ctx, param, value):
 
         ecid = value
         logger.debug('searching among connected devices via lockdownd')
-        for udid in list_devices():
+        for device in usbmux.list_devices():
             try:
-                lockdown = LockdownClient(udid=udid)
+                lockdown = LockdownClient(udid=device.serial)
             except IncorrectModeError:
                 continue
             if (ecid is None) or (lockdown.ecid == value):

pymobiledevice3/lockdown.py

@@ -1,14 +1,15 @@
 #!/usr/bin/env python3
+import datetime
 import logging
 import os
 import platform
 import plistlib
 import sys
 import uuid
-import datetime
 from pathlib import Path
 
 from packaging.version import Version
+
 from pymobiledevice3 import usbmux
 from pymobiledevice3.ca import ca_do_everything
 from pymobiledevice3.exceptions import *
@@ -33,24 +34,19 @@ def write_home_file(filename, data):
     return str(filepath)
 
 
-def list_devices():
-    return [d.serial for d in usbmux.list_devices()]
-
-
 class LockdownClient(object):
     DEFAULT_CLIENT_NAME = 'pyMobileDevice'
     SERVICE_PORT = 62078
 
     def __init__(self, udid=None, client_name=DEFAULT_CLIENT_NAME, autopair=True):
-        available_udids = list_devices()
-        if udid is None:
-            if len(available_udids) == 0:
-                raise NoDeviceConnectedError()
-            udid = available_udids[0]
-        else:
-            if (udid not in available_udids) and (udid.replace('-', '') not in available_udids):
+        device = usbmux.select_device(udid)
+        if device is None:
+            if udid:
                 raise ConnectionFailedError()
+            else:
+                raise NoDeviceConnectedError()
 
+        self.usbmux_device = device
         self.logger = logging.getLogger(__name__)
         self.paired = False
         self.SessionID = None
@@ -64,7 +60,7 @@ def __init__(self, udid=None, client_name=DEFAULT_CLIENT_NAME, autopair=True):
             raise IncorrectModeError()
 
         self.all_values = self.get_value()
-        self.udid = self.all_values.get('UniqueDeviceID', udid)
+        self.udid = self.all_values.get('UniqueDeviceID', self.usbmux_device.serial)
         self.unique_chip_id = self.all_values.get('UniqueChipID')
         self.device_public_key = self.all_values.get('DevicePublicKey')
         self.ios_version = self.all_values.get('ProductVersion')
@@ -206,10 +202,6 @@ def validate_pairing(self):
         return True
 
     def pair(self):
-        device_id = [
-            d for d in usbmux.list_devices()
-            if d.serial.replace('-', '') == self.udid.replace('-', '')
-        ][0].devid
         self.device_public_key = self.get_value('', 'DevicePublicKey')
         if not self.device_public_key:
             self.logger.error('Unable to retrieve DevicePublicKey')
@@ -243,7 +235,7 @@ def pair(self):
         record_data = plistlib.dumps(pair_record)
 
         client = usbmux.PlistProtocol(usbmux.MuxConnection.create_socket())
-        client.save_pair_record(self.udid, device_id, record_data)
+        client.save_pair_record(self.udid, self.usbmux_device.devid, record_data)
 
         self.paired = True
 

pymobiledevice3/restore/asr.py

@@ -6,8 +6,8 @@
 
 from tqdm import trange
 
+from pymobiledevice3 import usbmux
 from pymobiledevice3.exceptions import NoDeviceConnectedError, ConnectionFailedError, PyMobileDevice3Exception
-from pymobiledevice3.lockdown import list_devices
 from pymobiledevice3.service_connection import ServiceConnection
 
 ASR_VERSION = 1
@@ -26,18 +26,16 @@ class ASRClient(object):
     SERVICE_PORT = ASR_PORT
 
     def __init__(self, udid=None):
-        available_udids = list_devices()
-        if udid is None:
-            if len(available_udids) == 0:
-                raise NoDeviceConnectedError()
-            udid = available_udids[0]
-        else:
-            if udid not in available_udids:
+        device = usbmux.select_device(udid)
+        if device is None:
+            if udid:
                 raise ConnectionFailedError()
+            else:
+                raise NoDeviceConnectedError()
 
         logger.debug('connecting to ASR')
 
-        self.service = ServiceConnection.create(udid, self.SERVICE_PORT)
+        self.service = ServiceConnection.create(device.serial, self.SERVICE_PORT)
 
         logger.debug('ASR connected')
 
@@ -51,7 +49,7 @@ def __init__(self, udid=None):
 
         self.checksum_chunks = data.get('Checksum Chunks', False)
 
-    def recv_plist(self) -> dict:
+    def recv_plist(self) -> typing.Mapping:
         buf = b''
         while not buf.endswith(b'</plist>\n'):
             buf += self.service.recv()

pymobiledevice3/restore/fdr.py

@@ -5,8 +5,8 @@
 import threading
 from enum import Enum
 
+from pymobiledevice3 import usbmux
 from pymobiledevice3.exceptions import NoDeviceConnectedError, ConnectionFailedError, PyMobileDevice3Exception
-from pymobiledevice3.lockdown import list_devices
 from pymobiledevice3.service_connection import ServiceConnection
 
 CTRL_PORT = 0x43a  # 1082
@@ -36,22 +36,20 @@ class FDRClient:
     def __init__(self, type_: fdr_type, udid=None):
         global conn_port
 
-        available_udids = list_devices()
-        if udid is None:
-            if len(available_udids) == 0:
-                raise NoDeviceConnectedError()
-            udid = available_udids[0]
-        else:
-            if udid not in available_udids:
+        device = usbmux.select_device(udid)
+        if device is None:
+            if udid:
                 raise ConnectionFailedError()
+            else:
+                raise NoDeviceConnectedError()
 
         logger.debug('connecting to FDR')
 
         if type_ == fdr_type.FDR_CTRL:
-            self.service = ServiceConnection.create(udid, self.SERVICE_PORT)
+            self.service = ServiceConnection.create(device.serial, self.SERVICE_PORT)
             self.ctrl_handshake()
         else:
-            self.service = ServiceConnection.create(udid, conn_port)
+            self.service = ServiceConnection.create(device.serial, conn_port)
             self.sync_handshake()
 
         logger.debug('FDR connected')

pymobiledevice3/restore/restored_client.py

@@ -2,8 +2,8 @@
 
 from cached_property import cached_property
 
+from pymobiledevice3 import usbmux
 from pymobiledevice3.exceptions import NoDeviceConnectedError, ConnectionFailedError
-from pymobiledevice3.lockdown import list_devices
 from pymobiledevice3.restore.restore_options import RestoreOptions
 from pymobiledevice3.service_connection import ServiceConnection
 
@@ -24,15 +24,13 @@ def __init__(self, udid=None, client_name=DEFAULT_CLIENT_NAME):
 
     @staticmethod
     def _get_or_verify_udid(udid=None):
-        available_udids = list_devices()
-        if udid is None:
-            if len(available_udids) == 0:
-                raise NoDeviceConnectedError()
-            return available_udids[0]
-        else:
-            if udid not in available_udids:
+        device = usbmux.select_device(udid)
+        if device is None:
+            if udid:
                 raise ConnectionFailedError()
-            return udid
+            else:
+                raise NoDeviceConnectedError()
+        return device.serial
 
     def query_value(self, key=None):
         req = {'Request': 'QueryValue', 'Label': self.label}

pymobiledevice3/service_connection.py

@@ -9,6 +9,7 @@
 
 from pymobiledevice3 import usbmux
 from pymobiledevice3.exceptions import ConnectionFailedError, PyMobileDevice3Exception
+from pymobiledevice3.usbmux import select_device
 
 SHELL_USAGE = """
 # This shell allows you to communicate directly with every service layer behind the lockdownd daemon.
@@ -34,14 +35,9 @@ def __init__(self, socket):
 
     @staticmethod
     def create(udid, port):
-        target_device = None
+        target_device = select_device(udid)
         while target_device is None:
-            matching_devices = [device for device in usbmux.list_devices() if
-                                device.serial in (udid, udid.replace('-', ''))]
-            if len(matching_devices) == 1:
-                target_device = matching_devices[0]
-                break
-
+            target_device = select_device(udid)
         try:
             socket = target_device.connect(port)
         except usbmux.MuxException:

pymobiledevice3/services/debugserver_applist.py

@@ -1,5 +1,6 @@
 import logging
 import plistlib
+import typing
 
 from pymobiledevice3.lockdown import LockdownClient
 
@@ -14,7 +15,7 @@ def __init__(self, lockdown: LockdownClient):
         self.lockdown = lockdown
         self.service = self.lockdown.start_developer_service(self.SERVICE_NAME)
 
-    def get(self) -> dict:
+    def get(self) -> typing.Mapping:
         buf = b''
         while b'</plist>' not in buf:
             buf += self.service.recv(CHUNK_SIZE)

pymobiledevice3/usbmux.py

@@ -5,7 +5,7 @@
 import time
 from dataclasses import dataclass
 from enum import Enum
-from typing import List
+from typing import List, Optional
 
 from construct import Struct, Prefixed, Int32ul, GreedyBytes, StreamError, Int16ul, CString, Padding, FixedSized
 
@@ -204,7 +204,7 @@ def listen_for_devices(self, timeout=None):
             self.socket.sock.settimeout(end - time.time())
             try:
                 type_, data = self.proto.recv_device_state()
-                if type_ == PacketType.Attached:
+                if type_ == PacketType.Attached and data.is_legal:
                     self.devices.append(data)
                 elif type_ == PacketType.Detached:
                     self.devices = [device for device in self.devices if device.devid != data]
@@ -242,8 +242,25 @@ class MuxDevice:
     def connect(self, port) -> socket.socket:
         return create_mux().connect(self, port)
 
+    @property
+    def is_legal(self):
+        return bool(self.usbprod)
+
+    def matches_udid(self, udid):
+        return self.serial.replace('-', '') == udid.replace('-', '')
+
 
 def list_devices() -> List[MuxDevice]:
     mux = create_mux()
     mux.listen_for_devices(0.1)
     return mux.devices
+
+
+def select_device(udid='') -> Optional[MuxDevice]:
+    matching_devices = [
+        device for device in list_devices()
+        if device.is_legal and (not udid or device.matches_udid(udid))
+    ]
+    if not matching_devices:
+        return None
+    return matching_devices[0]