From 42331c3e37892f3ad2fc78307d9fe2b0ac292342 Mon Sep 17 00:00:00 2001
From: Richie B2B <richard@vdberg.org>
Date: Wed, 16 Feb 2022 23:45:11 +0100
Subject: [PATCH] Keyerror data token (#316)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
---
 pyicloud/services/drive.py | 42 ++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 11 deletions(-)

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."""