"""API for Google Drive bound to Home Assistant OAuth.""" from collections.abc import AsyncIterator, Callable, Coroutine from dataclasses import dataclass import json import logging from typing import Any, override from aiohttp import ClientSession, ClientTimeout, StreamReader from aiohttp.client_exceptions import ClientError, ClientResponseError from google_drive_api.api import AbstractAuth, GoogleDriveApi from homeassistant.components.backup import AgentBackup, suggested_filename from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_ACCESS_TOKEN from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import config_entry_oauth2_flow from .const import DOMAIN _UPLOAD_AND_DOWNLOAD_TIMEOUT = 12 * 3600 _UPLOAD_MAX_RETRIES = 20 _LOGGER = logging.getLogger(__name__) @dataclass class StorageQuotaData: """Class to represent storage quota data.""" limit: int | None usage: int usage_in_drive: int usage_in_trash: int class AsyncConfigEntryAuth(AbstractAuth): """Provide Google Drive authentication tied to an OAuth2 based config entry.""" def __init__( self, websession: ClientSession, oauth_session: config_entry_oauth2_flow.OAuth2Session, ) -> None: """Initialize AsyncConfigEntryAuth.""" super().__init__(websession) self._oauth_session = oauth_session @override async def async_get_access_token(self) -> str: """Return a valid access token.""" try: await self._oauth_session.async_ensure_token_valid() except ClientError as ex: is_setup = ( self._oauth_session.config_entry.state is ConfigEntryState.SETUP_IN_PROGRESS ) if isinstance(ex, ClientResponseError): # OAuth2 spec uses 400 for invalid refresh tokens. # During setup, we broaden this to all 4xx to abort cleanly. if ex.status == 400 or (is_setup and 400 <= ex.status < 500): if not is_setup: self._oauth_session.config_entry.async_start_reauth( self._oauth_session.hass ) raise ConfigEntryAuthFailed( translation_domain=DOMAIN, translation_key="authentication_not_valid", ) from ex if is_setup: raise ConfigEntryNotReady( translation_domain=DOMAIN, translation_key="authentication_failed", ) from ex raise return str(self._oauth_session.token[CONF_ACCESS_TOKEN]) class AsyncConfigFlowAuth(AbstractAuth): """Provide authentication tied to a fixed token for the config flow.""" def __init__( self, websession: ClientSession, token: str, ) -> None: """Initialize AsyncConfigFlowAuth.""" super().__init__(websession) self._token = token @override async def async_get_access_token(self) -> str: """Return a valid access token.""" return self._token class DriveClient: """Google Drive client.""" def __init__( self, ha_instance_id: str, auth: AbstractAuth, ) -> None: """Initialize Google Drive client.""" self._ha_instance_id = ha_instance_id self._api = GoogleDriveApi(auth) async def async_get_email_address(self) -> str: """Get email address of the current user.""" res = await self._api.get_user(params={"fields": "user(emailAddress)"}) return str(res["user"]["emailAddress"]) async def async_get_storage_quota(self) -> StorageQuotaData: """Get storage quota of the current user.""" res = await self._api.get_user(params={"fields": "storageQuota"}) storage_quota = res["storageQuota"] limit = storage_quota.get("limit") return StorageQuotaData( limit=int(limit) if limit is not None else None, usage=int(storage_quota.get("usage", 0)), usage_in_drive=int(storage_quota.get("usageInDrive", 0)), usage_in_trash=int(storage_quota.get("usageInTrash", 0)), ) async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]: """Create Home Assistant folder if it doesn't exist.""" fields = "id,name" query = " and ".join( [ "properties has { key='home_assistant' and value='root' }", "properties has { key='instance_id'" f" and value='{self._ha_instance_id}' }}", "trashed=false", ] ) res = await self._api.list_files( params={"q": query, "fields": f"files({fields})"} ) for file in res["files"]: _LOGGER.debug("Found existing folder: %s", file) return str(file["id"]), str(file["name"]) file_metadata = { "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", "properties": { "home_assistant": "root", "instance_id": self._ha_instance_id, }, } _LOGGER.debug("Creating new folder with metadata: %s", file_metadata) res = await self._api.create_file(params={"fields": fields}, json=file_metadata) _LOGGER.debug("Created folder: %s", res) return str(res["id"]), str(res["name"]) async def async_upload_backup( self, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, ) -> None: """Upload a backup.""" folder_id, _ = await self.async_create_ha_root_folder_if_not_exists() backup_metadata = { "name": suggested_filename(backup), "description": json.dumps(backup.as_dict()), "parents": [folder_id], "properties": { "home_assistant": "backup", "instance_id": self._ha_instance_id, "backup_id": backup.backup_id, }, } _LOGGER.debug( "Uploading backup: %s with Google Drive metadata: %s", backup.backup_id, backup_metadata, ) await self._api.resumable_upload_file( backup_metadata, open_stream, backup.size, max_retries=_UPLOAD_MAX_RETRIES, timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), ) _LOGGER.debug( "Uploaded backup: %s to: '%s'", backup.backup_id, backup_metadata["name"], ) async def async_list_backups(self) -> list[AgentBackup]: """List backups.""" query = " and ".join( [ "properties has { key='home_assistant' and value='backup' }", "properties has { key='instance_id'" f" and value='{self._ha_instance_id}' }}", "trashed=false", ] ) res = await self._api.list_files( params={"q": query, "fields": "files(description)"} ) backups = [] for file in res["files"]: backup = AgentBackup.from_dict(json.loads(file["description"])) backups.append(backup) return backups async def async_get_size_of_all_backups(self) -> int: """Get size of all backups.""" backups = await self.async_list_backups() return sum(backup.size for backup in backups) async def async_get_backup_file_id(self, backup_id: str) -> str | None: """Get file_id of backup if it exists.""" query = " and ".join( [ "properties has { key='home_assistant' and value='backup' }", "properties has { key='instance_id'" f" and value='{self._ha_instance_id}' }}", f"properties has {{ key='backup_id' and value='{backup_id}' }}", ] ) res = await self._api.list_files(params={"q": query, "fields": "files(id)"}) for file in res["files"]: return str(file["id"]) return None async def async_delete(self, file_id: str) -> None: """Delete file.""" await self._api.delete_file(file_id) async def async_download(self, file_id: str) -> StreamReader: """Download a file.""" resp = await self._api.get_file_content( file_id, timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT) ) return resp.content