import io import math import re from typing import Any, Dict, Optional from urllib.parse import urlencode from datetime import datetime, timedelta from aiohttp import ClientSession, ClientTimeout, ClientResponse from aiohttp.client_exceptions import ClientResponseError, ServerTimeoutError from injector import inject, singleton from ..util import AsyncHttpGetter from ..config import Config, Setting from ..exceptions import (GoogleCredentialsExpired, GoogleSessionError, LogicError, ProtocolError, ensureKey, KnownTransient, GoogleTimeoutError, GoogleUnexpectedError) from backup.util import Backoff from backup.file import JsonFileSaver from ..time import Time from ..logger import getLogger from backup.creds import Creds, Exchanger, DriveRequester from datetime import timezone logger = getLogger(__name__) MIME_TYPE = "application/tar" FOLDER_MIME_TYPE = 'application/vnd.google-apps.folder' FOLDER_NAME = 'Home Assistant Backups' DRIVE_VERSION = "v3" DRIVE_SERVICE = "drive" SELECT_FIELDS = "id,name,appProperties,size,trashed,mimeType,modifiedTime,capabilities,parents,driveId" THUMBNAIL_MIME_TYPE = "image/png" QUERY_FIELDS = "nextPageToken,files(" + SELECT_FIELDS + ")" CREATE_FIELDS = SELECT_FIELDS URL_FILES = "/drive/v3/files/" URL_ABOUT = "/drive/v3/about" URL_START_UPLOAD = "/upload/drive/v3/files/?uploadType=resumable&supportsAllDrives=true" PAGE_SIZE = 100 CHUNK_SIZE = 5 * 262144 RANGE_RE = re.compile("^bytes=0-\\d+$") BASE_CHUNK_SIZE = 256 * 1024 # Google's api requires uploading chunks in multiples of 256kb # During upload, chunks get sized to complete upload after 10s so we can give status updates on progress. CHUNK_UPLOAD_TARGET_SECONDS = 10 # don't attempt to resume a session with than this many times consistant failures, just in case something is broken on Google's # end so we don't retry the same broken session forever. Because the addon eventually backs off to doing 1 attempt/hour, this will # cause uploads to fail and start over after about 4 days. This gets reset every time a chunk successfully uploads. # God be with you if your upload takes that long. RETRY_SESSION_ATTEMPTS = 100 # Google claims that an upload session becomes invalid after 7 days. I have not verified this, but probably better to call it # after 6 and restart the session. UPLOAD_SESSION_EXPIRATION_DURATION = timedelta(days=6) RATE_LIMIT_EXCEEDED = 403 TOO_MANY_REQUESTS = 429 # Defines the retry strategy for calls made to Drive # max # of time to retry and call to Drive DRIVE_MAX_RETRIES: int = 5 # The initial backoff for drive retries. DRIVE_RETRY_INITIAL_SECONDS: int = 2 # How uch longer to wait for each Drive service call (Exponential backoff) DRIVE_EXPONENTIAL_BACKOFF: int = 2 OOB_CRED_CUTOFF = datetime(2022, 3, 16, tzinfo=timezone.utc) @singleton class DriveRequests(): @inject def __init__(self, config: Config, time: Time, drive: DriveRequester, session: ClientSession, exchanger: Exchanger): self.session = session self.config = config self.time = time self.drive = drive self.creds: Optional[Creds] = None self.exchanger: Exchanger = exchanger # Between attempts to upload, we keep track of the info needed to resume a resumable upload. self.last_attempt_metadata = None self.last_attempt_location = None self.last_attempt_count = 0 self.last_attempt_start_time = None self.tryLoadCredentials() async def _getHeaders(self): return { "Authorization": "Bearer " + await self.getToken(), "Client-Identifier": self.config.clientIdentifier() } @property def might_be_oob_creds(self): """Attempts to determine if the user might be using custom creds affected by google's OOB cred deprecation""" if not self.isCustomCreds(): return False if self.creds.original_expiration is None: # These creds must be old, so assume they're affected return True try: return self.creds.original_expiration < OOB_CRED_CUTOFF except: # noqa: E722 # Regardless of why this happens, assume they need to check return True def isCustomCreds(self): return self.creds is not None and self.creds.id != self.config.get(Setting.DEFAULT_DRIVE_CLIENT_ID) def _getAuthHeaders(self): return { "Client-Identifier": self.config.clientIdentifier() } def enabled(self): return self.creds is not None and self.config.get(Setting.ENABLE_DRIVE_UPLOAD) def _enabledCheck(self): if not self.enabled(): raise LogicError( "Attempt to use Google Drive before credentials are configured") def tryLoadCredentials(self): path = self.config.get(Setting.CREDENTIALS_FILE_PATH) if JsonFileSaver.exists(path): try: self.creds = Creds.load(self.time, JsonFileSaver.read(path)) except Exception: pass def saveCredentials(self, creds: Creds): path = self.config.get(Setting.CREDENTIALS_FILE_PATH) if not creds: if JsonFileSaver.exists(path): JsonFileSaver.delete(path) self.creds = None return JsonFileSaver.write(path, creds.serialize()) self.tryLoadCredentials() async def getToken(self, refresh=False): if self.creds and not self.creds.is_expired and not refresh: return self.creds.access_token # refresh the credentials logger.debug("Requesting refreshed Google Drive credentials") self.creds = await self.exchanger.refresh(self.creds) return self.creds.access_token async def refreshToken(self): await self.getToken(refresh=True) async def get(self, id): q = { "fields": SELECT_FIELDS, "supportsAllDrives": "true" } async with await self.retryRequest("GET", URL_FILES + id + "/?" + urlencode(q)) as response: return await response.json() async def download(self, id, size): ret = AsyncHttpGetter(self.config.get(Setting.DRIVE_URL) + URL_FILES + id + "/?alt=media&supportsAllDrives=true", await self._getHeaders(), self.session, size=size, timeoutFactory=GoogleTimeoutError.factory, otherErrorFactory=GoogleUnexpectedError.factory, timeout=ClientTimeout( sock_connect=self.config.get(Setting.DOWNLOAD_TIMEOUT_SECONDS), sock_read=self.config.get(Setting.DOWNLOAD_TIMEOUT_SECONDS)), time=self.time) return ret async def query(self, query): # SOMEDAY: Add a test for page size, test server support is needed too for continuation tokens continuation = None while True: q = { "q": query, "fields": QUERY_FIELDS, "pageSize": self.config.get(Setting.GOOGLE_DRIVE_PAGE_SIZE), "supportsAllDrives": "true", "includeItemsFromAllDrives": "true", "corpora": "allDrives" } if continuation: q["pageToken"] = continuation async with await self.retryRequest("GET", URL_FILES + "?" + urlencode(q)) as response: data = await response.json() for item in data['files']: yield item if "nextPageToken" not in data or len(data['nextPageToken']) <= 0: break else: continuation = data['nextPageToken'] async def update(self, id, update_metadata): async with await self.retryRequest("PATCH", URL_FILES + id + "/?supportsAllDrives=true", json=update_metadata): pass async def delete(self, id): async with await self.retryRequest("DELETE", URL_FILES + id + "/?supportsAllDrives=true"): pass async def getAboutInfo(self): q = {"fields": 'storageQuota,user'} async with await self.retryRequest("GET", URL_ABOUT + "?" + urlencode(q)) as resp: return await resp.json() async def create(self, stream, metadata, mime_type): # Upload logic is complicated. See https://developers.google.com/drive/api/v3/manage-uploads#resumable total_size = stream.size() location = None if metadata == self.last_attempt_metadata and self.last_attempt_location is not None and self.last_attempt_count < RETRY_SESSION_ATTEMPTS and self.time.now() < self.last_attempt_start_time + UPLOAD_SESSION_EXPIRATION_DURATION: logger.debug( "Attempting to resume a previously failed upload where we left off") self.last_attempt_count += 1 # Attempt to resume from a partially completed upload. headers = { "Content-Length": "0", "Content-Range": "bytes */{0}".format(total_size) } try: async with await self.retryRequest("PUT", self.last_attempt_location, headers=headers, patch_url=False) as initial: if initial.status == 308: # We can resume the upload, check where it left off if 'Range' in initial.headers: position = int(initial.headers["Range"][len("bytes=0-"):]) stream.position(position + 1) else: # No range header in the response means no bytes have been uploaded yet. stream.position(0) logger.debug("Resuming upload at byte {0} of {1}".format( stream.position(), total_size)) location = self.last_attempt_location else: logger.debug("Drive returned status code {0}, so we'll have to start the upload over again.".format( initial.status)) except ClientResponseError as e: if e.status == 410: # Drive doesn't recognize the resume token, so we'll just have to start over. logger.debug("Drive upload session wasn't recognized, restarting upload from the beginning.") location = None else: raise if location is None: # There is no session resume, so start a new one. logger.debug("Starting a new upload session with Google Drive") headers = { "X-Upload-Content-Type": mime_type, "X-Upload-Content-Length": str(total_size), } async with await self.retryRequest("POST", URL_START_UPLOAD, headers=headers, json=metadata) as initial: # Google returns a url in the header "Location", which is where subsequent requests to upload # the backup's bytes should be sent. Logic below handles uploading the file bytes in chunks. location = ensureKey( 'Location', initial.headers, "Google Drive's Upload headers") self.last_attempt_count = 0 stream.position(0) # Keep track of the location in case the upload fails and we want to resume where we left off. # "metadata" is a durable fingerprint that uniquely identifies a backup, so we can use it to identify a # resumable partial upload in future retrys. self.last_attempt_location = location self.last_attempt_metadata = metadata self.last_attempt_start_time = self.time.now() # Always start with the minimum chunk size and work up from there in case the last attempt # failed due to connectivity errors or ... whatever. current_chunk_size = BASE_CHUNK_SIZE while True: start = stream.position() data = await stream.read(current_chunk_size) chunk_size = len(data.getbuffer()) if chunk_size == 0: raise LogicError( "Backup file stream ended prematurely while uploading to Google Drive") headers = { "Content-Length": str(chunk_size), "Content-Range": "bytes {0}-{1}/{2}".format(start, start + chunk_size - 1, total_size) } startTime = self.time.now() logger.debug("Sending {0} bytes to Google Drive".format(current_chunk_size)) try: async with await self.retryRequest("PUT", location, headers=headers, data=data, patch_url=False) as partial: # Base the next chunk size on how long it took to send the last chunk. current_chunk_size = self._getNextChunkSize( current_chunk_size, (self.time.now() - startTime).total_seconds()) # any time a chunk gets uploaded, reset the retry counter. This lets very flaky connections # complete eventually after enough retrying. self.last_attempt_count = 1 yield float(start + chunk_size) / float(total_size) if partial.status == 200 or partial.status == 201: # Upload completed, return the object json self.last_attempt_location = None self.last_attempt_metadata = None yield await self.get((await partial.json())['id']) break elif partial.status == 308: # Upload partially complete, seek to the new requested position range_bytes = ensureKey( "Range", partial.headers, "Google Drive's upload response headers") if not RANGE_RE.match(range_bytes): raise ProtocolError( "Range", partial.headers, "Google Drive's upload response headers") position = int(partial.headers["Range"][len("bytes=0-"):]) stream.position(position + 1) else: partial.raise_for_status() except ClientResponseError as e: if math.floor(e.status / 100) == 4: # clear the cached session location URI, since a 4XX error # always means the upload session is no good anymore (AFAIK) self.last_attempt_location = None self.last_attempt_metadata = None if e.status == 404: raise GoogleSessionError() else: raise e def _getNextChunkSize(self, last_chunk_size, last_chunk_seconds): max = BASE_CHUNK_SIZE * math.floor(self.config.get(Setting.MAXIMUM_UPLOAD_CHUNK_BYTES) / BASE_CHUNK_SIZE) if max < BASE_CHUNK_SIZE: max = BASE_CHUNK_SIZE if last_chunk_seconds <= 0: return max next_chunk = CHUNK_UPLOAD_TARGET_SECONDS * last_chunk_size / last_chunk_seconds if next_chunk >= max: return max if next_chunk < BASE_CHUNK_SIZE: return BASE_CHUNK_SIZE return math.floor(next_chunk / BASE_CHUNK_SIZE) * BASE_CHUNK_SIZE async def createFolder(self, metadata): async with await self.retryRequest("POST", URL_FILES + "?supportsAllDrives=true", json=metadata) as resp: return await resp.json() async def retryRequest(self, method, url, auth_headers: Optional[Dict[str, str]] = None, headers: Optional[Dict[str, str]] = None, json: Optional[Dict[str, Any]] = None, data: Any = None, cred_retry: bool = True, patch_url: bool = True) -> ClientResponse: backoff = Backoff(base=DRIVE_RETRY_INITIAL_SECONDS, attempts=DRIVE_MAX_RETRIES) if patch_url: url = self.config.get(Setting.DRIVE_URL) + url while True: headers_to_use = await self._getHeaders() if headers: headers_to_use.update(headers) if self.config.get(Setting.TRACE_REQUESTS): logger.trace("Making Google Drive request: " + url) try: data_to_use = data if isinstance(data_to_use, io.BytesIO): # This is a pretty low-down dirty hack, but it works and lets us reuse the byte stream. # aiohttp complains if you pass it a large byte object data_to_use = io.BytesIO(data_to_use.getbuffer()) data_to_use.seek(0) return await self.drive.request(method, url, headers=headers_to_use, json=json, data=data_to_use) except GoogleCredentialsExpired: # Get fresh credentials, then retry right away. logger.debug("Google Drive credentials have expired. We'll retry with new ones.") await self.refreshToken() except KnownTransient as e: backoff.backoff(e) logger.error("{0}: we'll retry in {1} seconds".format(e.message(), backoff.peek())) await self.time.sleepAsync(backoff.peek()) except ServerTimeoutError: raise GoogleTimeoutError()