Spaces:
Sleeping
Sleeping
File size: 5,649 Bytes
80cb919 e3a165a 80cb919 e3a165a 80cb919 e3a165a 80cb919 e3a165a 80cb919 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
# GCS credential token refresher
import os, json, logging
from typing import Optional
# Dynamic imports for Google OAuth (only when not in local mode)
def _import_google_oauth():
"""Dynamically import Google OAuth libraries only when needed"""
try:
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import Flow
from google.auth.transport.requests import Request
return Credentials, Flow, Request
except ImportError as e:
raise ImportError(f"Google OAuth libraries not available: {e}. Make sure IS_LOCAL=false and google-auth packages are installed.")
# Check if we're in local mode
IS_LOCAL = os.getenv("IS_LOCAL", "false").lower() == "true"
# Only import Google OAuth libraries if not in local mode
if not IS_LOCAL:
try:
Credentials, Flow, Request = _import_google_oauth()
except ImportError:
# Create dummy classes for when Google OAuth is not available
class Credentials:
@staticmethod
def from_authorized_user_info(*args, **kwargs):
raise ImportError("Google OAuth not available")
class Flow:
@staticmethod
def from_client_config(*args, **kwargs):
raise ImportError("Google OAuth not available")
class Request:
pass
else:
# Create dummy classes for local mode
class Credentials:
@staticmethod
def from_authorized_user_info(*args, **kwargs):
raise ImportError("Google OAuth not available in local mode")
class Flow:
@staticmethod
def from_client_config(*args, **kwargs):
raise ImportError("Google OAuth not available in local mode")
class Request:
pass
logger = logging.getLogger("token")
if not logger.handlers:
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
logger.addHandler(handler)
SCOPES = ["https://www.googleapis.com/auth/drive.file"]
TOKEN_FILE = os.getenv("GDRIVE_TOKEN_FILE", "cache/secrets/gdrive_token.json")
def _load_oauth_client_web():
cfg_env = os.getenv("GDRIVE_CREDENTIALS_JSON")
if not cfg_env:
return None
try:
cfg = json.loads(cfg_env)
return cfg.get("web")
except Exception as e:
logger.error(f"❌ Failed to parse GDRIVE_CREDENTIALS_JSON: {e}")
return None
def _ensure_dirs():
base = os.path.dirname(TOKEN_FILE)
if base and not os.path.exists(base):
os.makedirs(base, exist_ok=True)
def get_credentials() -> Optional[Credentials]:
if IS_LOCAL:
logger.info("🏠 Local mode: Google OAuth credentials not needed")
return None
# 1) Token file
if os.path.exists(TOKEN_FILE):
try:
with open(TOKEN_FILE, "r", encoding="utf-8") as f:
data = json.load(f)
creds = Credentials.from_authorized_user_info(data, scopes=SCOPES)
if creds and creds.expired and creds.refresh_token:
creds.refresh(Request())
logger.info("🔄 Refreshed access token from token file")
return creds
except Exception as e:
logger.warning(f"⚠️ Failed to load token file: {e}")
# 2) Refresh token in env
refresh = os.getenv("GDRIVE_REFRESH_TOKEN")
web = _load_oauth_client_web()
if refresh and web:
creds = Credentials(
None,
refresh_token=refresh,
token_uri="https://oauth2.googleapis.com/token",
client_id=web.get("client_id"),
client_secret=web.get("client_secret"),
scopes=SCOPES,
)
if creds and (creds.expired or not creds.valid):
try:
creds.refresh(Request())
logger.info("🔄 Refreshed access token from env refresh token")
except Exception as e:
logger.warning(f"⚠️ Refresh with env token failed: {e}")
return creds
# 3) Nothing available
return None
def build_auth_url(redirect_uri: str) -> str:
if IS_LOCAL:
raise RuntimeError("Google OAuth not available in local mode")
web = _load_oauth_client_web()
if not web:
raise RuntimeError("GDRIVE_CREDENTIALS_JSON missing or invalid ('web' section required)")
flow = Flow.from_client_config({"web": web}, scopes=SCOPES, redirect_uri=redirect_uri)
auth_url, _ = flow.authorization_url(
prompt="consent",
access_type="offline",
include_granted_scopes="true"
)
return auth_url
def exchange_code(code: str, redirect_uri: str) -> Credentials:
if IS_LOCAL:
raise RuntimeError("Google OAuth not available in local mode")
web = _load_oauth_client_web()
if not web:
raise RuntimeError("GDRIVE_CREDENTIALS_JSON missing or invalid ('web' section required)")
flow = Flow.from_client_config({"web": web}, scopes=SCOPES, redirect_uri=redirect_uri)
flow.fetch_token(code=code)
creds: Credentials = flow.credentials
info = {
"token": creds.token,
"refresh_token": creds.refresh_token,
"token_uri": "https://oauth2.googleapis.com/token",
"client_id": web.get("client_id"),
"client_secret": web.get("client_secret"),
"scopes": SCOPES,
}
_ensure_dirs()
with open(TOKEN_FILE, "w", encoding="utf-8") as f:
json.dump(info, f)
logger.info("✅ Saved Google refresh token to %s", TOKEN_FILE)
# also set env for current process
if creds.refresh_token:
os.environ["GDRIVE_REFRESH_TOKEN"] = creds.refresh_token
return creds |