266 lines
No EOL
14 KiB
Python
266 lines
No EOL
14 KiB
Python
"""
|
|
Token Manager Service
|
|
Handles all token operations including automatic refresh for backend services.
|
|
"""
|
|
|
|
import logging
|
|
import httpx
|
|
from typing import Optional, Dict, Any, Callable
|
|
|
|
from modules.datamodels.datamodelSecurity import Token
|
|
from modules.datamodels.datamodelUam import AuthAuthority
|
|
from modules.shared.configuration import APP_CONFIG
|
|
from modules.shared.timezoneUtils import get_utc_timestamp, create_expiration_timestamp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class TokenManager:
|
|
"""Centralized token management service"""
|
|
|
|
def __init__(self):
|
|
# Microsoft OAuth configuration
|
|
self.msft_client_id = APP_CONFIG.get("Service_MSFT_CLIENT_ID")
|
|
self.msft_client_secret = APP_CONFIG.get("Service_MSFT_CLIENT_SECRET")
|
|
self.msft_tenant_id = APP_CONFIG.get("Service_MSFT_TENANT_ID", "common")
|
|
|
|
# Google OAuth configuration
|
|
self.google_client_id = APP_CONFIG.get("Service_GOOGLE_CLIENT_ID")
|
|
self.google_client_secret = APP_CONFIG.get("Service_GOOGLE_CLIENT_SECRET")
|
|
|
|
def refresh_microsoft_token(self, refresh_token: str, user_id: str, old_token: Token) -> Optional[Token]:
|
|
"""Refresh Microsoft OAuth token using refresh token"""
|
|
try:
|
|
logger.debug(f"refresh_microsoft_token: Starting Microsoft token refresh for user {user_id}")
|
|
logger.debug(f"refresh_microsoft_token: Configuration check - client_id: {bool(self.msft_client_id)}, client_secret: {bool(self.msft_client_secret)}")
|
|
|
|
if not self.msft_client_id or not self.msft_client_secret:
|
|
logger.error("Microsoft OAuth configuration not found")
|
|
return None
|
|
|
|
# Microsoft token refresh endpoint
|
|
token_url = f"https://login.microsoftonline.com/{self.msft_tenant_id}/oauth2/v2.0/token"
|
|
logger.debug(f"refresh_microsoft_token: Using token URL: {token_url}")
|
|
|
|
# Prepare refresh request
|
|
data = {
|
|
"client_id": self.msft_client_id,
|
|
"client_secret": self.msft_client_secret,
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": refresh_token,
|
|
"scope": "Mail.ReadWrite Mail.Send Mail.ReadWrite.Shared User.Read"
|
|
}
|
|
logger.debug(f"refresh_microsoft_token: Refresh request data prepared (refresh_token length: {len(refresh_token) if refresh_token else 0})")
|
|
|
|
# Make refresh request
|
|
with httpx.Client(timeout=30.0) as client:
|
|
logger.debug(f"refresh_microsoft_token: Making HTTP request to Microsoft OAuth endpoint")
|
|
response = client.post(token_url, data=data)
|
|
logger.debug(f"refresh_microsoft_token: HTTP response status: {response.status_code}")
|
|
|
|
if response.status_code == 200:
|
|
token_data = response.json()
|
|
logger.debug(f"refresh_microsoft_token: Token refresh successful, creating new token")
|
|
|
|
# Create new token
|
|
new_token = Token(
|
|
userId=user_id,
|
|
authority=AuthAuthority.MSFT,
|
|
connectionId=old_token.connectionId, # Preserve connection ID
|
|
tokenAccess=token_data["access_token"],
|
|
tokenRefresh=token_data.get("refresh_token", refresh_token), # Keep old refresh token if new one not provided
|
|
tokenType=token_data.get("token_type", "bearer"),
|
|
expiresAt=create_expiration_timestamp(token_data.get("expires_in", 3600)),
|
|
createdAt=get_utc_timestamp()
|
|
)
|
|
|
|
logger.debug(f"refresh_microsoft_token: New token created with ID: {new_token.id}")
|
|
return new_token
|
|
else:
|
|
logger.error(f"Failed to refresh Microsoft token: {response.status_code} - {response.text}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error refreshing Microsoft token: {str(e)}")
|
|
return None
|
|
|
|
def refresh_google_token(self, refresh_token: str, user_id: str, old_token: Token) -> Optional[Token]:
|
|
"""Refresh Google OAuth token using refresh token"""
|
|
try:
|
|
logger.debug(f"refresh_google_token: Starting Google token refresh for user {user_id}")
|
|
logger.debug(f"refresh_google_token: Configuration check - client_id: {bool(self.google_client_id)}, client_secret: {bool(self.google_client_secret)}")
|
|
|
|
if not self.google_client_id or not self.google_client_secret:
|
|
logger.error("Google OAuth configuration not found")
|
|
return None
|
|
|
|
# Google token refresh endpoint
|
|
token_url = "https://oauth2.googleapis.com/token"
|
|
logger.debug(f"refresh_google_token: Using token URL: {token_url}")
|
|
|
|
# Prepare refresh request
|
|
data = {
|
|
"client_id": self.google_client_id,
|
|
"client_secret": self.google_client_secret,
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": refresh_token
|
|
}
|
|
logger.debug(f"refresh_google_token: Refresh request data prepared (refresh_token length: {len(refresh_token) if refresh_token else 0})")
|
|
|
|
# Make refresh request
|
|
with httpx.Client(timeout=30.0) as client:
|
|
logger.debug(f"refresh_google_token: Making HTTP request to Google OAuth endpoint")
|
|
response = client.post(token_url, data=data)
|
|
logger.debug(f"refresh_google_token: HTTP response status: {response.status_code}")
|
|
|
|
if response.status_code == 200:
|
|
token_data = response.json()
|
|
logger.debug(f"refresh_google_token: Token refresh successful, creating new token")
|
|
|
|
# Validate the response contains required fields
|
|
if "access_token" not in token_data:
|
|
logger.error("Google token refresh response missing access_token")
|
|
return None
|
|
|
|
# Create new token
|
|
new_token = Token(
|
|
userId=user_id,
|
|
authority=AuthAuthority.GOOGLE,
|
|
connectionId=old_token.connectionId, # Preserve connection ID
|
|
tokenAccess=token_data["access_token"],
|
|
tokenRefresh=token_data.get("refresh_token", refresh_token), # Use new refresh token if provided
|
|
tokenType=token_data.get("token_type", "bearer"),
|
|
expiresAt=create_expiration_timestamp(token_data.get("expires_in", 3600)),
|
|
createdAt=get_utc_timestamp()
|
|
)
|
|
|
|
logger.debug(f"refresh_google_token: New token created with ID: {new_token.id}")
|
|
return new_token
|
|
else:
|
|
error_details = response.text
|
|
logger.error(f"Failed to refresh Google token: {response.status_code} - {error_details}")
|
|
|
|
# Handle specific error cases
|
|
if response.status_code == 400:
|
|
try:
|
|
error_data = response.json()
|
|
error_code = error_data.get("error")
|
|
if error_code == "invalid_grant":
|
|
logger.warning("Google refresh token is invalid or expired - user needs to re-authenticate")
|
|
elif error_code == "invalid_client":
|
|
logger.error("Google OAuth client configuration is invalid")
|
|
except:
|
|
pass
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error refreshing Google token: {str(e)}")
|
|
return None
|
|
|
|
def refresh_token(self, old_token: Token) -> Optional[Token]:
|
|
"""Refresh an expired token using the appropriate OAuth service"""
|
|
try:
|
|
logger.debug(f"refresh_token: Starting refresh for token {old_token.id}, authority: {old_token.authority}")
|
|
logger.debug(f"refresh_token: Token details: userId={old_token.userId}, connectionId={old_token.connectionId}, hasRefreshToken={bool(old_token.tokenRefresh)}")
|
|
|
|
# Cooldown: avoid refreshing too frequently if a workflow triggers refresh repeatedly
|
|
# Only allow a new refresh if at least 10 minutes passed since the token was created/refreshed
|
|
try:
|
|
now_ts = get_utc_timestamp()
|
|
created_ts = float(old_token.createdAt) if old_token.createdAt is not None else 0.0
|
|
seconds_since_last_refresh = now_ts - created_ts
|
|
if seconds_since_last_refresh < 10 * 60:
|
|
logger.info(
|
|
f"refresh_token: Skipping refresh for connection {old_token.connectionId} due to cooldown. "
|
|
f"Last refresh {int(seconds_since_last_refresh)}s ago (< 600s)."
|
|
)
|
|
# Return the existing token to avoid caller errors while preventing provider rate limits
|
|
return old_token
|
|
except Exception:
|
|
# If any issue reading timestamps, proceed with normal refresh to be safe
|
|
pass
|
|
|
|
if not old_token.tokenRefresh:
|
|
logger.warning(f"No refresh token available for {old_token.authority}")
|
|
return None
|
|
|
|
# Route to appropriate refresh method
|
|
if old_token.authority == AuthAuthority.MSFT:
|
|
logger.debug(f"refresh_token: Refreshing Microsoft token")
|
|
return self.refresh_microsoft_token(old_token.tokenRefresh, old_token.userId, old_token)
|
|
elif old_token.authority == AuthAuthority.GOOGLE:
|
|
logger.debug(f"refresh_token: Refreshing Google token")
|
|
return self.refresh_google_token(old_token.tokenRefresh, old_token.userId, old_token)
|
|
else:
|
|
logger.warning(f"Unknown authority for token refresh: {old_token.authority}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error refreshing token: {str(e)}")
|
|
return None
|
|
|
|
def ensure_fresh_token(self, token: Token, *, seconds_before_expiry: int = 30 * 60, save_callback: Optional[Callable[[Token], None]] = None) -> Optional[Token]:
|
|
"""Ensure a token is fresh; refresh if expiring within threshold.
|
|
|
|
Args:
|
|
token: Existing token to validate/refresh.
|
|
seconds_before_expiry: Threshold window to proactively refresh.
|
|
save_callback: Optional function to persist a refreshed token.
|
|
|
|
Returns:
|
|
A fresh token (refreshed or original) or None if refresh failed.
|
|
"""
|
|
try:
|
|
if token is None:
|
|
return None
|
|
|
|
now_ts = get_utc_timestamp()
|
|
expires_at = token.expiresAt or 0
|
|
|
|
# If token expires within the threshold, try to refresh
|
|
if expires_at and expires_at < (now_ts + seconds_before_expiry):
|
|
logger.info(
|
|
f"ensure_fresh_token: Token for connection {token.connectionId} expiring soon "
|
|
f"(in {max(0, expires_at - now_ts)}s). Attempting proactive refresh."
|
|
)
|
|
refreshed = self.refresh_token(token)
|
|
if refreshed:
|
|
if save_callback is not None:
|
|
try:
|
|
save_callback(refreshed)
|
|
except Exception as e:
|
|
logger.warning(f"ensure_fresh_token: Failed to persist refreshed token: {e}")
|
|
return refreshed
|
|
else:
|
|
logger.warning("ensure_fresh_token: Token refresh failed")
|
|
return None
|
|
|
|
# Token is sufficiently fresh
|
|
return token
|
|
except Exception as e:
|
|
logger.error(f"ensure_fresh_token: Error ensuring fresh token: {e}")
|
|
return None
|
|
|
|
# Convenience wrapper to fetch and ensure fresh token for a connection via interface layer
|
|
def getFreshToken(self, connectionId: str, secondsBeforeExpiry: int = 30 * 60) -> Optional[Token]:
|
|
"""Return a fresh token for a connection, refreshing when expiring soon.
|
|
|
|
Reads the latest stored token via interface layer, then
|
|
uses ensure_fresh_token to refresh if needed and persists the refreshed
|
|
token via interface layer.
|
|
"""
|
|
try:
|
|
from modules.interfaces.interfaceDbAppObjects import getRootInterface
|
|
interfaceDbApp = getRootInterface()
|
|
|
|
token = interfaceDbApp.getConnectionToken(connectionId)
|
|
if not token:
|
|
return None
|
|
return self.ensure_fresh_token(
|
|
token,
|
|
seconds_before_expiry=secondsBeforeExpiry,
|
|
save_callback=lambda t: interfaceDbApp.saveConnectionToken(t)
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"getFreshToken: Error fetching or refreshing token for connection {connectionId}: {e}")
|
|
return None |