"""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 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, HomeAssistantError, ) 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 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: if ( self._oauth_session.config_entry.state is ConfigEntryState.SETUP_IN_PROGRESS ): if isinstance(ex, ClientResponseError) and 400 <= ex.status < 500: raise ConfigEntryAuthFailed( translation_domain=DOMAIN, translation_key="authentication_not_valid", ) from ex raise ConfigEntryNotReady( translation_domain=DOMAIN, translation_key="authentication_failed", ) from ex if hasattr(ex, "status") and ex.status == 400: self._oauth_session.config_entry.async_start_reauth( self._oauth_session.hass ) raise HomeAssistantError( translation_domain=DOMAIN, translation_key="authentication_failed", ) from ex 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 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