New Addon
Google BK
This commit is contained in:
@@ -0,0 +1,373 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user