diff --git a/pyicloud/services/drive.py b/pyicloud/services/drive.py index 7c11b74..87573ce 100644 --- a/pyicloud/services/drive.py +++ b/pyicloud/services/drive.py @@ -1,6 +1,7 @@ """Drive service.""" from datetime import datetime, timedelta import json +import logging import io import mimetypes import os @@ -8,6 +9,11 @@ import time from re import search from requests import Response +from pyicloud.exceptions import PyiCloudAPIResponseException + + +LOGGER = logging.getLogger(__name__) + class DriveService: """The 'Drive' iCloud service.""" @@ -42,6 +48,7 @@ class DriveService: ] ), ) + self._raise_if_error(request) return request.json()[0] def get_file(self, file_id, **kwargs): @@ -52,16 +59,22 @@ class DriveService: self._document_root + "/ws/com.apple.CloudDocs/download/by_id", params=file_params, ) - if not response.ok: - return None - url = response.json()["data_token"]["url"] - return self.session.get(url, params=self.params, **kwargs) + self._raise_if_error(response) + response_json = response.json() + package_token = response_json.get("package_token") + data_token = response_json.get("data_token") + if data_token and data_token.get("url"): + return self.session.get(data_token["url"], params=self.params, **kwargs) + if package_token and package_token.get("url"): + return self.session.get(package_token["url"], params=self.params, **kwargs) + raise KeyError("'data_token' nor 'package_token'") def get_app_data(self): """Returns the app library (previously ubiquity).""" request = self.session.get( self._service_root + "/retrieveAppLibraries", params=self.params ) + self._raise_if_error(request) return request.json()["items"] def _get_upload_contentws_url(self, file_object): @@ -92,8 +105,7 @@ class DriveService: } ), ) - if not request.ok: - return None + self._raise_if_error(request) return (request.json()[0]["document_id"], request.json()[0]["url"]) def _update_contentws(self, folder_id, sf_info, document_id, file_object): @@ -131,8 +143,7 @@ class DriveService: headers={"Content-Type": "text/plain"}, data=json.dumps(data), ) - if not request.ok: - return None + self._raise_if_error(request) return request.json() def send_file(self, folder_id, file_object): @@ -140,10 +151,8 @@ class DriveService: document_id, content_url = self._get_upload_contentws_url(file_object) request = self.session.post(content_url, files={file_object.name: file_object}) - if not request.ok: - return None + self._raise_if_error(request) content_response = request.json()["singleFile"] - self._update_contentws(folder_id, content_response, document_id, file_object) def create_folders(self, parent, name): @@ -164,6 +173,7 @@ class DriveService: } ), ) + self._raise_if_error(request) return request.json() def rename_items(self, node_id, etag, name): @@ -183,6 +193,7 @@ class DriveService: } ), ) + self._raise_if_error(request) return request.json() def move_items_to_trash(self, node_id, etag): @@ -202,6 +213,7 @@ class DriveService: } ), ) + self._raise_if_error(request) return request.json() @property @@ -217,6 +229,14 @@ class DriveService: def __getitem__(self, key): return self.root[key] + def _raise_if_error(self, response): # pylint: disable=no-self-use + if not response.ok: + api_error = PyiCloudAPIResponseException( + response.reason, response.status_code + ) + LOGGER.error(api_error) + raise api_error + class DriveNode: """Drive node."""