Fix DeviceLocked raise on lookup image action

apptim · May 8, 2024 · 095d3e5 · 095d3e5
1 parent ad4c4af
commit 095d3e5
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/pymobiledevice3/services/mobile_image_mounter.py b/pymobiledevice3/services/mobile_image_mounter.py
@@ -48,6 +48,8 @@ def lookup_image(self, image_type: str) -> bytes:
         response = self.service.send_recv_plist({'Command': 'LookupImage',
                                                  'ImageType': image_type})
 
+        if not response or 'DeviceLocked' in response.get('Error', ''):
+            raise DeviceLockedError()
         if not response or not response.get('ImagePresent', True):
             raise NotMountedError()
 
@@ -116,10 +118,7 @@ def upload_image(self, image_type: str, image: bytes, signature: bytes) -> None:
         status = result.get('Status')
 
         if status != 'ReceiveBytesAck':
-            if 'DeviceLocked' in result.get('Error', ''):
-                raise DeviceLockedError()
-            else:
-                raise PyMobileDevice3Exception(f'command ReceiveBytes failed with: {result}')
+            raise PyMobileDevice3Exception(f'command ReceiveBytes failed with: {result}')
 
         self.service.sendall(image)
         result = self.service.recv_plist()