Coverage for icloudpy/services/drive.py: 100%
183 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 15:10 +0000
1"""Drive service."""
3import io
4import json
5import mimetypes
6import os
7import time
8from datetime import datetime, timedelta
9from re import search
11from requests import Response
14class DriveService:
15 """The 'Drive' iCloud service."""
17 def __init__(self, service_root, document_root, session, params):
18 self._service_root = service_root
19 self._document_root = document_root
20 self.session = session
21 self.params = dict(params)
22 self._root = None
24 def _get_token_from_cookie(self):
25 for cookie in self.session.cookies:
26 if cookie.name == "X-APPLE-WEBAUTH-VALIDATE":
27 match = search(r"\bt=([^:]+)", cookie.value)
28 if match is None:
29 raise Exception(f"Can't extract token from {cookie.value}")
30 return {"token": match.group(1)}
31 raise Exception("Token cookie not found")
33 def get_node_data(self, drivewsid):
34 """Returns the node data."""
35 request = self.session.post(
36 self._service_root + "/retrieveItemDetailsInFolders",
37 params=self.params,
38 data=json.dumps(
39 [
40 {
41 "drivewsid": drivewsid,
42 "partialData": False,
43 },
44 ],
45 ),
46 )
47 if not request.ok:
48 self.session.raise_error(request.status_code, request.reason)
49 return request.json()[0]
51 def get_file(self, file_id, zone="com.apple.CloudDocs", **kwargs):
52 """Returns iCloud Drive file."""
53 file_params = dict(self.params)
54 file_params.update({"document_id": file_id})
55 # The lookup and the transfer are separate requests. A timeout given
56 # only to the transfer leaves the lookup free to stall its thread
57 # forever, so forward it here too. ``stream`` is deliberately not
58 # forwarded: the lookup is a small JSON reply read in full.
59 lookup_kwargs = {"timeout": kwargs["timeout"]} if "timeout" in kwargs else {}
60 response = self.session.get(
61 self._document_root + f"/ws/{zone}/download/by_id",
62 params=file_params,
63 **lookup_kwargs,
64 )
65 if not response.ok:
66 self.session.raise_error(response.status_code, response.reason)
67 package_token = response.json().get("package_token")
68 data_token = response.json().get("data_token")
69 if data_token and data_token.get("url"):
70 return self.session.get(data_token["url"], params=self.params, **kwargs)
71 elif package_token and package_token.get("url"):
72 return self.session.get(package_token["url"], params=self.params, **kwargs)
73 else:
74 raise KeyError("'data_token' nor 'package_token' found in response.")
76 def get_app_data(self):
77 """Returns the app library."""
78 request = self.session.get(
79 self._service_root + "/retrieveAppLibraries",
80 params=self.params,
81 )
82 if not request.ok:
83 self.session.raise_error(request.status_code, request.reason)
84 return request.json()["items"]
86 def get_app_node(self, app_id, folder="documents"):
87 """Returns the node of the app."""
88 return DriveNode(self, self.get_node_data("FOLDER::" + app_id + "::" + folder))
90 def _get_upload_contentws_url(self, file_object, zone="com.apple.CloudDocs"):
91 """Get the contentWS endpoint URL to add a new file."""
92 content_type = mimetypes.guess_type(file_object.name)[0]
93 if content_type is None:
94 content_type = ""
96 # Get filesize from file object
97 orig_pos = file_object.tell()
98 file_object.seek(0, os.SEEK_END)
99 file_size = file_object.tell()
100 file_object.seek(orig_pos, os.SEEK_SET)
102 file_params = self.params
103 file_params.update(self._get_token_from_cookie())
105 request = self.session.post(
106 self._document_root + f"/ws/{zone}/upload/web",
107 params=file_params,
108 headers={"Content-Type": "text/plain"},
109 data=json.dumps(
110 {
111 "filename": file_object.name,
112 "type": "FILE",
113 "content_type": content_type,
114 "size": file_size,
115 },
116 ),
117 )
118 if not request.ok:
119 self.session.raise_error(request.status_code, request.reason)
120 return (request.json()[0]["document_id"], request.json()[0]["url"])
122 def _update_contentws(
123 self,
124 folder_id,
125 sf_info,
126 document_id,
127 file_object,
128 zone="com.apple.CloudDocs",
129 ):
130 data = {
131 "data": {
132 "signature": sf_info["fileChecksum"],
133 "wrapping_key": sf_info["wrappingKey"],
134 "reference_signature": sf_info["referenceChecksum"],
135 "size": sf_info["size"],
136 },
137 "command": "add_file",
138 "create_short_guid": True,
139 "document_id": document_id,
140 "path": {
141 "starting_document_id": folder_id,
142 "path": os.path.basename(file_object.name),
143 },
144 "allow_conflict": True,
145 "file_flags": {
146 "is_writable": True,
147 "is_executable": False,
148 "is_hidden": False,
149 },
150 "mtime": int(time.time() * 1000),
151 "btime": int(time.time() * 1000),
152 }
154 # Add the receipt if we have one. Will be absent for 0-sized files
155 if sf_info.get("receipt"):
156 data["data"].update({"receipt": sf_info["receipt"]})
158 request = self.session.post(
159 self._document_root + f"/ws/{zone}/update/documents",
160 params=self.params,
161 headers={"Content-Type": "text/plain"},
162 data=json.dumps(data),
163 )
164 if not request.ok:
165 self.session.raise_error(request.status_code, request.reason)
166 return request.json()
168 def send_file(self, folder_id, file_object, zone="com.apple.CloudDocs"):
169 """Send new file to iCloud Drive."""
170 document_id, content_url = self._get_upload_contentws_url(file_object, zone)
172 request = self.session.post(content_url, files={file_object.name: file_object})
173 if not request.ok:
174 self.session.raise_error(request.status_code, request.reason)
175 content_response = request.json()["singleFile"]
177 self._update_contentws(
178 folder_id,
179 content_response,
180 document_id,
181 file_object,
182 zone,
183 )
185 def create_folders(self, parent, name):
186 """Creates a new iCloud Drive folder"""
187 request = self.session.post(
188 self._service_root + "/createFolders",
189 params=self.params,
190 headers={"Content-Type": "text/plain"},
191 data=json.dumps(
192 {
193 "destinationDrivewsId": parent,
194 "folders": [
195 {
196 "clientId": self.params["clientId"],
197 "name": name,
198 },
199 ],
200 },
201 ),
202 )
203 return request.json()
205 def rename_items(self, node_id, etag, name):
206 """Renames an iCloud Drive node"""
207 request = self.session.post(
208 self._service_root + "/renameItems",
209 params=self.params,
210 data=json.dumps(
211 {
212 "items": [
213 {
214 "drivewsid": node_id,
215 "etag": etag,
216 "name": name,
217 },
218 ],
219 },
220 ),
221 )
222 return request.json()
224 def move_items_to_trash(self, node_id, etag):
225 """Moves an iCloud Drive node to the trash bin"""
226 request = self.session.post(
227 self._service_root + "/moveItemsToTrash",
228 params=self.params,
229 data=json.dumps(
230 {
231 "items": [
232 {
233 "drivewsid": node_id,
234 "etag": etag,
235 "clientId": self.params["clientId"],
236 },
237 ],
238 },
239 ),
240 )
241 if not request.ok:
242 self.session.raise_error(request.status_code, request.reason)
243 return request.json()
245 @property
246 def root(self):
247 """Returns the root node."""
248 if not self._root:
249 self._root = DriveNode(
250 self,
251 self.get_node_data("FOLDER::com.apple.CloudDocs::root"),
252 )
253 return self._root
255 def __getattr__(self, attr):
256 return getattr(self.root, attr)
258 def __getitem__(self, key):
259 return self.root[key]
262class DriveNode:
263 """Drive node."""
265 def __init__(self, conn, data):
266 self.data = data
267 self.connection = conn
268 self._children = None
270 @property
271 def name(self):
272 """Gets the node name."""
273 if "extension" in self.data:
274 return f"{self.data['name']}.{self.data['extension']}"
275 return self.data["name"]
277 @property
278 def type(self):
279 """Gets the node type."""
280 node_type = self.data.get("type")
281 return node_type and node_type.lower()
283 def get_children(self):
284 """Gets the node children."""
285 if not self._children:
286 if "items" not in self.data:
287 self.data.update(self.connection.get_node_data(self.data["drivewsid"]))
288 if "items" not in self.data:
289 raise KeyError(f"No items in folder, status: {self.data['status']}")
290 self._children = [DriveNode(self.connection, item_data) for item_data in self.data["items"]]
291 return self._children
293 @property
294 def size(self):
295 """Gets the node size."""
296 size = self.data.get("size") # Folder does not have size
297 if not size:
298 return None
299 return int(size)
301 @property
302 def date_created(self):
303 """Gets the node created date (in UTC)."""
304 return _date_to_utc(self.data.get("dateCreated"))
306 @property
307 def date_changed(self):
308 """Gets the node changed date (in UTC)."""
309 return _date_to_utc(self.data.get("dateChanged")) # Folder does not have date
311 @property
312 def date_modified(self):
313 """Gets the node modified date (in UTC)."""
314 return _date_to_utc(self.data.get("dateModified")) # Folder does not have date
316 @property
317 def date_last_open(self):
318 """Gets the node last open date (in UTC)."""
319 return _date_to_utc(self.data.get("lastOpenTime")) # Folder does not have date
321 def open(self, **kwargs):
322 """Gets the node file."""
323 # iCloud returns 400 Bad Request for 0-byte files
324 if self.data["size"] == 0:
325 response = Response()
326 response.raw = io.BytesIO()
327 return response
328 return self.connection.get_file(
329 self.data["docwsid"],
330 zone=self.data["zone"],
331 **kwargs,
332 )
334 def upload(self, file_object, **kwargs):
335 """ "Upload a new file."""
336 return self.connection.send_file(
337 self.data["docwsid"],
338 file_object,
339 zone=self.data["zone"],
340 **kwargs,
341 )
343 def dir(self):
344 """Gets the node list of directories."""
345 if self.type == "file":
346 return None
347 return [child.name for child in self.get_children()]
349 def mkdir(self, folder):
350 """Create a new directory directory."""
351 # remove cached entries information first so that it will be re-read on next get_children()
352 self._children = None
353 if "items" in self.data:
354 self.data.pop("items")
355 return self.connection.create_folders(self.data["drivewsid"], folder)
357 def rename(self, name):
358 """Rename an iCloud Drive item."""
359 return self.connection.rename_items(
360 self.data["drivewsid"],
361 self.data["etag"],
362 name,
363 )
365 def delete(self):
366 """Delete an iCloud Drive item."""
367 return self.connection.move_items_to_trash(
368 self.data["drivewsid"],
369 self.data["etag"],
370 )
372 def get(self, name):
373 """Gets the node child."""
374 if self.type == "file":
375 return None
376 return [child for child in self.get_children() if child.name == name][0]
378 def __getitem__(self, key):
379 try:
380 return self.get(key)
381 except IndexError as error:
382 raise KeyError(f"No child named '{key}' exists") from error
384 def __unicode__(self):
385 return f"{{type: {self.type}, name: {self.name}}}"
387 def __str__(self):
388 return self.__unicode__()
390 def __repr__(self):
391 return f"<{type(self).__name__}: {str(self)}>"
394def _date_to_utc(date):
395 if not date:
396 return None
397 # jump through hoops to return time in UTC rather than California time
398 match = search(r"^(.+?)([\+\-]\d+):(\d\d)$", date)
399 if not match:
400 # Already in UTC
401 return datetime.strptime(date, "%Y-%m-%dT%H:%M:%SZ")
402 base = datetime.strptime(match.group(1), "%Y-%m-%dT%H:%M:%S")
403 diff = timedelta(hours=int(match.group(2)), minutes=int(match.group(3)))
404 return base - diff