work

main.py

@@ -16,79 +16,10 @@ from pymobiledevice3.remote.tunnel_service import CoreDeviceTunnelProxy, start_t
 from pymobiledevice3.service_connection import ServiceConnection
 from pymobiledevice3.services.dvt.dvt_secure_socket_proxy import DvtSecureSocketProxyService
 from pymobiledevice3.services.dvt.instruments.location_simulation import LocationSimulation
+from pymobiledevice3.services.mobile_image_mounter import MobileImageMounterService
 from pymobiledevice3.usbmux import list_devices
-
-_orig_remotexpc_connect = RemoteXPCConnection.connect
-
-
-async def _remotexpc_connect_with_timeout(self):
-    """Wrapper to add timeout to TCP connection and handshake"""
-    import logging
-    _logger = logging.getLogger("ios-api")
-
-    _logger.info(f"RemoteXPC attempting TCP connection to {self.address}")
-    tcp_start = time.time()
-    try:
-        self._reader, self._writer = await asyncio.wait_for(
-            asyncio.open_connection(self.address[0], self.address[1]),
-            timeout=5.0
-        )
-        tcp_elapsed = time.time() - tcp_start
-        _logger.info(f"RemoteXPC TCP connected in {tcp_elapsed:.2f}s")
-    except asyncio.TimeoutError as exc:
-        tcp_elapsed = time.time() - tcp_start
-        _logger.error(f"RemoteXPC TCP connection to {self.address} timed out after {tcp_elapsed:.2f}s")
-        raise asyncio.TimeoutError(f"TCP connection to {self.address} timed out after 5s") from exc
-    except Exception as exc:
-        tcp_elapsed = time.time() - tcp_start
-        _logger.error(f"RemoteXPC TCP connection to {self.address} failed after {tcp_elapsed:.2f}s: {exc}")
-        raise
-
-    _logger.info(f"RemoteXPC starting handshake")
-    handshake_start = time.time()
-    try:
-        await self._do_handshake()
-        handshake_elapsed = time.time() - handshake_start
-        _logger.info(f"RemoteXPC handshake complete in {handshake_elapsed:.2f}s")
-    except Exception as exc:
-        handshake_elapsed = time.time() - handshake_start
-        _logger.error(f"RemoteXPC handshake failed after {handshake_elapsed:.2f}s: {exc}")
-        await self.close()
-        raise
-
-
-RemoteXPCConnection.connect = _remotexpc_connect_with_timeout
-
 from pymobiledevice3.remote.remote_service_discovery import RemoteServiceDiscoveryService
 
-# Patch RemoteServiceDiscoveryService.connect to use faster timeouts for lockdown
-_orig_rsd_connect = RemoteServiceDiscoveryService.connect
-
-
-async def _rsd_connect_with_timeout(self):
-    """Modified connect with faster timeout for lockdown connections"""
-    import logging
-
-    _logger = logging.getLogger("ios-api")
-
-    await self.service.connect()
-    try:
-        self.peer_info = await self.service.receive_response()
-        self.udid = self.peer_info["Properties"]["UniqueDeviceID"]
-        self.product_type = self.peer_info["Properties"]["ProductType"]
-
-        # Skip lockdown connection - not working over RSD tunnel and not needed for DVT
-        self.lockdown = None
-        _logger.info("Skipping lockdown connection for RSD (not required for DVT)")
-
-        self.all_values = self.lockdown.all_values if self.lockdown is not None else {}
-    except Exception:
-        await self.close()
-        raise
-
-
-RemoteServiceDiscoveryService.connect = _rsd_connect_with_timeout
-
 
 class JsonFormatter(logging.Formatter):
     def format(self, record: logging.LogRecord) -> str:
@@ -103,65 +34,6 @@ class JsonFormatter(logging.Formatter):
         return json.dumps(payload, ensure_ascii=True)
 
 
-handler = logging.StreamHandler()
-handler.setFormatter(JsonFormatter())
-root_logger = logging.getLogger()
-root_logger.handlers = [handler]
-root_logger.setLevel(logging.INFO)
-
-logger = logging.getLogger("ios-api")
-
-# Patch ServiceConnection.create_using_tcp to force IPv6 and reduce timeout
-_orig_create_using_tcp = ServiceConnection.create_using_tcp
-
-
-def _create_using_tcp_with_ipv6(
-        hostname: str,
-        port: int,
-        keep_alive: bool = True,
-        create_connection_timeout: int = 5,  # Reduced from 20 to 5 seconds
-):
-    """Force IPv6 connection with reduced timeout"""
-    import socket as socket_module
-    from pymobiledevice3.osu.os_utils import get_os_utils
-    import logging
-    _logger = logging.getLogger("ios-api")
-
-    # Force IPv6 socket creation for tunnel addresses
-    if ':' in hostname:  # IPv6 address
-        _logger.info(f"ServiceConnection connecting to {hostname}:{port} with {create_connection_timeout}s timeout")
-        sock = socket_module.socket(socket_module.AF_INET6, socket_module.SOCK_STREAM)
-        sock.settimeout(create_connection_timeout)
-        connect_start = time.time()
-        try:
-            sock.connect((hostname, port))
-            connect_elapsed = time.time() - connect_start
-            _logger.info(f"ServiceConnection connected to {hostname}:{port} in {connect_elapsed:.2f}s")
-            # Keep a 5-second timeout for subsequent operations for faster failure detection
-            sock.settimeout(5.0)
-            if keep_alive:
-                get_os_utils().set_keepalive(sock)
-            return ServiceConnection(sock)
-        except Exception as exc:
-            connect_elapsed = time.time() - connect_start
-            _logger.error(f"ServiceConnection failed to {hostname}:{port} after {connect_elapsed:.2f}s: {exc}")
-            sock.close()
-            raise
-    else:
-        # Fall back to original for non-IPv6
-        return _orig_create_using_tcp(
-            hostname,
-            port,
-            keep_alive=keep_alive,
-            create_connection_timeout=create_connection_timeout,
-        )
-
-
-ServiceConnection.create_using_tcp = staticmethod(_create_using_tcp_with_ipv6)
-
-app = FastAPI(title="iOS Device Management API")
-
-
 class LocationUpdate(BaseModel):
     latitude: float
     longitude: float
@@ -169,34 +41,28 @@ class LocationUpdate(BaseModel):
     rsd_port: Optional[int] = None
 
 
-@app.get("/api/usbmux/status")
-async def get_usbmux_status():
-    """Lists all devices visible to USBMux."""
-    try:
-        device = list_devices()
-        return {
-            "device_count": len(device),
-            "devices": [{"serial": d.serial, "connection_type": d.connection_type} for d in device]
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+class DeviceShort(BaseModel):
+    udid: str
+    connection_type: str
 
 
-@app.get("/api/lockdown/status")
-async def get_lockdown_status():
-    """Checks lockdown connectivity and basic device info."""
-    try:
-        device_serial = _get_single_device_udid()
-        lockdown = create_using_usbmux(serial=device_serial, autopair=False)
-        return {
-            "udid": device_serial,
-            "product_version": lockdown.product_version,
-            "device_name": lockdown.get_value(key='DeviceName'),
-            "phone_number": lockdown.get_value(key='PhoneNumber'),
-            "status": "Connected"
-        }
-    except Exception as e:
-        return {"status": "Disconnected", "error": str(e)}
+class DeviceStatus(BaseModel):
+    device_connected: bool = False
+    device_count: int = 0
+    devices: Optional[list[DeviceShort]] = None
+    udid: Optional[str] = None
+    device_name: Optional[str] = None
+    product_version: Optional[str] = None
+    phone_number: Optional[str] = None
+    developer_mode_enabled: Optional[bool] = None
+    ddi_mounted: Optional[bool] = None
+    rsd_address: Optional[str] = None
+    rsd_port: Optional[int] = None
+    lockdown_trusted_port: Optional[int] = None
+    lockdown_untrusted_port: Optional[int] = None
+    lockdown_trusted_reachable: bool = False
+    lockdown_untrusted_reachable: bool = False
+    dtservicehub_reachable: bool = False
 
 
 class TunnelStartRequest(BaseModel):
@@ -230,6 +96,143 @@ class TunnelState:
 
 _TUNNELS: Dict[str, TunnelState] = {}
 
+_orig_rsd_connect = RemoteServiceDiscoveryService.connect
+_orig_remotexpc_connect = RemoteXPCConnection.connect
+_orig_create_using_tcp = ServiceConnection.create_using_tcp
+
+# Global reference to keep the CLI process alive
+_location_sim_process: Optional[asyncio.subprocess.Process] = None
+
+# Global DVT session for reuse
+_dvt_session: Optional[DvtSecureSocketProxyService] = None
+_dvt_session_lock = asyncio.Lock()
+_current_rsd: Optional[RemoteServiceDiscoveryService] = None
+
+handler = logging.StreamHandler()
+handler.setFormatter(JsonFormatter())
+root_logger = logging.getLogger()
+root_logger.handlers = [handler]
+root_logger.setLevel(logging.INFO)
+logger = logging.getLogger("ios-api")
+
+
+async def _remotexpc_connect_with_timeout(self):
+    """Wrapper to add timeout to TCP connection and handshake"""
+    logger.info(f"RemoteXPC attempting TCP connection to {self.address}")
+    tcp_start = time.time()
+    try:
+        self._reader, self._writer = await asyncio.wait_for(
+            asyncio.open_connection(self.address[0], self.address[1]),
+            timeout=5.0
+        )
+        tcp_elapsed = time.time() - tcp_start
+        logger.info(f"RemoteXPC TCP connected in {tcp_elapsed:.2f}s")
+    except asyncio.TimeoutError as exc:
+        tcp_elapsed = time.time() - tcp_start
+        logger.error(f"RemoteXPC TCP connection to {self.address} timed out after {tcp_elapsed:.2f}s")
+        raise asyncio.TimeoutError(f"TCP connection to {self.address} timed out after 5s") from exc
+    except Exception as exc:
+        tcp_elapsed = time.time() - tcp_start
+        logger.error(f"RemoteXPC TCP connection to {self.address} failed after {tcp_elapsed:.2f}s: {exc}")
+        raise
+
+    logger.info(f"RemoteXPC starting handshake")
+    handshake_start = time.time()
+    try:
+        await self._do_handshake()
+        handshake_elapsed = time.time() - handshake_start
+        logger.info(f"RemoteXPC handshake complete in {handshake_elapsed:.2f}s")
+    except Exception as exc:
+        handshake_elapsed = time.time() - handshake_start
+        logger.error(f"RemoteXPC handshake failed after {handshake_elapsed:.2f}s: {exc}")
+        await self.close()
+        raise
+
+
+async def _rsd_connect_with_timeout(self):
+    """Modified connect with faster timeout for lockdown connections"""
+    await self.service.connect()
+    try:
+        self.peer_info = await self.service.receive_response()
+        self.udid = self.peer_info["Properties"]["UniqueDeviceID"]
+        self.product_type = self.peer_info["Properties"]["ProductType"]
+
+        # Skip lockdown connection - not working over RSD tunnel and not needed for DVT
+        self.lockdown = None
+        logger.info("Skipping lockdown connection for RSD (not required for DVT)")
+
+        self.all_values = self.lockdown.all_values if self.lockdown is not None else {}
+    except Exception:
+        await self.close()
+        raise
+
+
+RemoteXPCConnection.connect = _remotexpc_connect_with_timeout
+RemoteServiceDiscoveryService.connect = _rsd_connect_with_timeout
+
+
+def _create_using_tcp_with_ipv6(
+        hostname: str,
+        port: int,
+        keep_alive: bool = True,
+        create_connection_timeout: int = 5,  # Reduced from 20 to 5 seconds
+):
+    """Force IPv6 connection with reduced timeout"""
+    import socket as socket_module
+    from pymobiledevice3.osu.os_utils import get_os_utils
+
+    # Force IPv6 socket creation for tunnel addresses
+    if ':' in hostname:  # IPv6 address
+        logger.info(f"ServiceConnection connecting to {hostname}:{port} with {create_connection_timeout}s timeout")
+        sock = socket_module.socket(socket_module.AF_INET6, socket_module.SOCK_STREAM)
+        sock.settimeout(create_connection_timeout)
+        connect_start = time.time()
+        try:
+            sock.connect((hostname, port))
+            connect_elapsed = time.time() - connect_start
+            logger.info(f"ServiceConnection connected to {hostname}:{port} in {connect_elapsed:.2f}s")
+            # Keep a 5-second timeout for subsequent operations for faster failure detection
+            sock.settimeout(5.0)
+            if keep_alive:
+                get_os_utils().set_keepalive(sock)
+            return ServiceConnection(sock)
+        except Exception as exc:
+            connect_elapsed = time.time() - connect_start
+            logger.error(f"ServiceConnection failed to {hostname}:{port} after {connect_elapsed:.2f}s: {exc}")
+            sock.close()
+            raise
+    else:
+        # Fall back to original for non-IPv6
+        return _orig_create_using_tcp(
+            hostname,
+            port,
+            keep_alive=keep_alive,
+            create_connection_timeout=create_connection_timeout,
+        )
+
+
+ServiceConnection.create_using_tcp = staticmethod(_create_using_tcp_with_ipv6)
+
+
+def _get_developer_mode_status() -> bool:
+    device_serial = _get_single_device_udid()
+    lockdown = create_using_usbmux(serial=device_serial, autopair=False)
+    return True if MobileImageMounterService(lockdown).query_developer_mode_status() == 1 else False
+
+def _get_developer_disk_image_status() -> bool:
+    device_serial = _get_single_device_udid()
+    lockdown = create_using_usbmux(serial=device_serial, autopair=False)
+    images = MobileImageMounterService(lockdown).copy_devices()
+    is_ddi = False
+    for image in images:
+        if image.get("DiskImageType") == "Personalized" and image.get("PersonalizedImageType") == "DeveloperDiskImage":
+            is_ddi = True
+
+    return is_ddi
+
+
+
+    return True if MobileImageMounterService(lockdown).is_image_mounted() == 1 else False
 
 def _get_single_device_udid() -> str:
     devices = list_devices()
@@ -396,15 +399,6 @@ def _simulate_location_with_dvt(service_provider, latitude: float, longitude: fl
             dvt.service.close()
 
 
-# Global reference to keep the CLI process alive
-_location_sim_process: Optional[asyncio.subprocess.Process] = None
-
-# Global DVT session for reuse
-_dvt_session: Optional[DvtSecureSocketProxyService] = None
-_dvt_session_lock = asyncio.Lock()
-_current_rsd: Optional[RemoteServiceDiscoveryService] = None
-
-
 async def _get_or_create_dvt_session(rsd: RemoteServiceDiscoveryService):
     """Get existing DVT session or create a new one (with proper initialization)"""
     global _dvt_session, _current_rsd
@@ -687,41 +681,74 @@ async def _preflight_rsd_async(
         await rsd.close()
 
 
+app = FastAPI(title="iOS Device Management API")
+
+
+@app.get("/api/status")
+async def get_status():
+    """Lists all devices visible to USBMux."""
+    try:
+        devices = list_devices()
+        if len(devices) > 1:
+            device_list = []
+            for d in devices:
+                device_list.append(DeviceShort(udid=d.serial, connection_type=d.connection_type))
+            return DeviceStatus(device_connected=True, device_count=len(devices), devices=device_list)
+        try:
+            lockdown = create_using_usbmux(serial=devices[0].serial, autopair=False)
+            state = _TUNNELS.get("")
+            if state and state.result:
+                rsd_address = state.result.address
+                rsd_port = state.result.port
+            else:
+                udid = devices[0].serial
+                logger.info("Auto-starting tunnel")
+                result = await _start_tunnel_internal(
+                    udid,
+                    TunnelProtocol.TCP,
+                    True,
+                    30,
+                )
+                rsd_address = result.address
+                rsd_port = result.port
+
+            return DeviceStatus(device_connected=True, device_count=len(devices), udid=devices[0].serial,
+                                device_name=lockdown.get_value(key='DeviceName'),
+                                product_version=lockdown.product_version,
+                                phone_number=lockdown.get_value(key='PhoneNumber'),
+                                developer_mode_enabled=_get_developer_mode_status(),
+                                ddi_mounted=_get_developer_disk_image_status(),
+                                rsd_address=rsd_address, rsd_port=rsd_port, )
+        except Exception as e:
+            logger.info("Error establishing lockdown: %s", str(e))
+            return DeviceStatus(device_connected=False, device_count=0)
+    except Exception as e:
+        logger.info("No device connected: %s", str(e))
+        return DeviceStatus(device_connected=False, device_count=0)
+
+
+@app.get("/api/lockdown/status")
+async def get_lockdown_status():
+    """Checks lockdown connectivity and basic device info."""
+    try:
+        device_serial = _get_single_device_udid()
+        lockdown = create_using_usbmux(serial=device_serial, autopair=False)
+        return {
+            "udid": device_serial,
+            "product_version": lockdown.product_version,
+            "device_name": lockdown.get_value(key='DeviceName'),
+            "phone_number": lockdown.get_value(key='PhoneNumber'),
+            "status": "Connected"
+        }
+    except Exception as e:
+        return {"status": "Disconnected", "error": str(e)}
+
+
 @app.post("/api/tunnel/start")
 async def start_usb_tunnel(data: TunnelStartRequest):
     """Starts a CoreDevice tunnel to a USB device and returns RSD connection details."""
     udid = _get_single_device_udid()
     key = ""
-    existing = _TUNNELS.get(key)
-    if existing and not existing.task.done():
-        if existing.result is not None:
-            logger.info("Tunnel already running udid=%s", existing.udid)
-            return {
-                "status": "already_running",
-                "udid": existing.udid or udid,
-                "rsd_address": existing.result.address,
-                "rsd_port": existing.result.port,
-                "interface": existing.result.interface,
-                "protocol": existing.result.protocol.name.lower(),
-            }
-
-    protocol = _parse_protocol(data.protocol)
-    result = await _start_tunnel_internal(
-        udid,
-        protocol,
-        data.wait_for_device,
-        data.wait_timeout_seconds,
-    )
-
-    logger.info("Tunnel start completed udid=%s", udid)
-    return {
-        "status": "started",
-        "udid": udid,
-        "rsd_address": result.address,
-        "rsd_port": result.port,
-        "interface": result.interface,
-        "protocol": result.protocol.name.lower(),
-    }
 
 
 @app.post("/api/tunnel/stop")