136 lines
3.9 KiB
Python
136 lines
3.9 KiB
Python
"""Google OAuth client with Google Drive API support.
|
|
|
|
Activated when OAUTH_GOOGLE_ID is set in .env.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import httpx
|
|
|
|
from app.core.config import get_settings
|
|
|
|
GOOGLE_OAUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
|
|
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
|
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"
|
|
GOOGLE_DRIVE_URL = "https://www.googleapis.com/drive/v3"
|
|
|
|
SCOPES = " ".join([
|
|
"https://www.googleapis.com/auth/userinfo.email",
|
|
"https://www.googleapis.com/auth/userinfo.profile",
|
|
"https://www.googleapis.com/auth/drive.file",
|
|
])
|
|
|
|
VOIDEA_DRIVE_FOLDER = "VoIdeaAI"
|
|
|
|
|
|
@dataclass
|
|
class GoogleUserInfo:
|
|
id: str
|
|
email: str
|
|
display_name: str
|
|
avatar_url: str | None
|
|
|
|
|
|
def is_available() -> bool:
|
|
return get_settings().google_oauth_enabled
|
|
|
|
|
|
async def get_authorize_url() -> str:
|
|
settings = get_settings()
|
|
params = {
|
|
"response_type": "code",
|
|
"client_id": settings.oauth_google_id,
|
|
"redirect_uri": settings.oauth_google_redirect_uri,
|
|
"scope": SCOPES,
|
|
"access_type": "offline",
|
|
"prompt": "consent",
|
|
}
|
|
query = "&".join(f"{k}={v}" for k, v in params.items())
|
|
return f"{GOOGLE_OAUTH_URL}?{query}"
|
|
|
|
|
|
async def exchange_code(code: str) -> dict | None:
|
|
settings = get_settings()
|
|
if not settings.google_oauth_enabled:
|
|
return None
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(
|
|
GOOGLE_TOKEN_URL,
|
|
data={
|
|
"grant_type": "authorization_code",
|
|
"code": code,
|
|
"client_id": settings.oauth_google_id,
|
|
"client_secret": settings.oauth_google_secret,
|
|
"redirect_uri": settings.oauth_google_redirect_uri,
|
|
},
|
|
)
|
|
if resp.status_code != 200:
|
|
return None
|
|
return resp.json()
|
|
|
|
|
|
async def get_user_info(access_token: str) -> GoogleUserInfo | None:
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get(
|
|
GOOGLE_USERINFO_URL,
|
|
headers={"Authorization": f"Bearer {access_token}"},
|
|
)
|
|
if resp.status_code != 200:
|
|
return None
|
|
data = resp.json()
|
|
return GoogleUserInfo(
|
|
id=data["id"],
|
|
email=data["email"],
|
|
display_name=data.get("name", data["email"]),
|
|
avatar_url=data.get("picture"),
|
|
)
|
|
|
|
|
|
async def ensure_app_folder(access_token: str) -> str | None:
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(
|
|
f"{GOOGLE_DRIVE_URL}/files",
|
|
headers={
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"name": VOIDEA_DRIVE_FOLDER,
|
|
"mimeType": "application/vnd.google-apps.folder",
|
|
},
|
|
)
|
|
if resp.status_code == 200:
|
|
return resp.json().get("id")
|
|
return None
|
|
|
|
|
|
async def upload_file(
|
|
access_token: str,
|
|
file_name: str,
|
|
file_content: bytes,
|
|
parent_id: str | None = None,
|
|
) -> bool:
|
|
metadata = {"name": file_name}
|
|
if parent_id:
|
|
metadata["parents"] = [parent_id]
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(
|
|
f"{GOOGLE_DRIVE_URL}/files?uploadType=multipart",
|
|
headers={"Authorization": f"Bearer {access_token}"},
|
|
data={"metadata": str(metadata)},
|
|
files={"file": (file_name, file_content)},
|
|
)
|
|
return resp.status_code == 200
|
|
|
|
|
|
async def get_disk_info(access_token: str) -> dict | None:
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get(
|
|
f"{GOOGLE_DRIVE_URL}/about?fields=storageQuota",
|
|
headers={"Authorization": f"Bearer {access_token}"},
|
|
)
|
|
if resp.status_code != 200:
|
|
return None
|
|
return resp.json()
|