diff --git a/pymobiledevice3/services/mobile_image_mounter.py b/pymobiledevice3/services/mobile_image_mounter.py index c73c0cb73..3903e0382 100755 --- 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()