import base64 import hashlib import json import logging import time import uuid from typing import List, Optional from cryptography.fernet import Fernet from rexpro_ai.env import OAUTH_SESSION_TOKEN_ENCRYPTION_KEY from rexpro_ai.internal.db import Base, get_async_db_context from pydantic import BaseModel, ConfigDict from sqlalchemy import BigInteger, Column, Index, String, Text, delete, select, update from sqlalchemy.ext.asyncio import AsyncSession log = logging.getLogger(__name__) #################### # DB MODEL #################### class OAuthSession(Base): __tablename__ = 'oauth_session' id = Column(Text, primary_key=True, unique=True) user_id = Column(Text, nullable=False) provider = Column(Text, nullable=False) token = Column(Text, nullable=False) # JSON with access_token, id_token, refresh_token expires_at = Column(BigInteger, nullable=False) created_at = Column(BigInteger, nullable=False) updated_at = Column(BigInteger, nullable=False) # Add indexes for better performance __table_args__ = ( Index('idx_oauth_session_user_id', 'user_id'), Index('idx_oauth_session_expires_at', 'expires_at'), Index('idx_oauth_session_user_provider', 'user_id', 'provider'), ) class OAuthSessionModel(BaseModel): id: str user_id: str provider: str token: dict expires_at: int # timestamp in epoch created_at: int # timestamp in epoch updated_at: int # timestamp in epoch model_config = ConfigDict(from_attributes=True) #################### # Forms #################### class OAuthSessionResponse(BaseModel): id: str user_id: str provider: str expires_at: int class OAuthSessionTable: def __init__(self): self.encryption_key = OAUTH_SESSION_TOKEN_ENCRYPTION_KEY if not self.encryption_key: raise Exception('OAUTH_SESSION_TOKEN_ENCRYPTION_KEY is not set') # check if encryption key is in the right format for Fernet (32 url-safe base64-encoded bytes) if len(self.encryption_key) != 44: key_bytes = hashlib.sha256(self.encryption_key.encode()).digest() self.encryption_key = base64.urlsafe_b64encode(key_bytes) else: self.encryption_key = self.encryption_key.encode() try: self.fernet = Fernet(self.encryption_key) except Exception as e: log.error(f'Error initializing Fernet with provided key: {e}') raise def _encrypt_token(self, token) -> str: """Encrypt OAuth tokens for storage""" try: token_json = json.dumps(token) encrypted = self.fernet.encrypt(token_json.encode()).decode() return encrypted except Exception as e: log.error(f'Error encrypting tokens: {e}') raise def _decrypt_token(self, token: str): """Decrypt OAuth tokens from storage""" try: decrypted = self.fernet.decrypt(token.encode()).decode() return json.loads(decrypted) except Exception as e: log.error(f'Error decrypting tokens: {type(e).__name__}: {e}') raise async def create_session( self, user_id: str, provider: str, token: dict, db: Optional[AsyncSession] = None, ) -> Optional[OAuthSessionModel]: """Create a new OAuth session""" try: async with get_async_db_context(db) as db: current_time = int(time.time()) id = str(uuid.uuid4()) result = OAuthSession( **{ 'id': id, 'user_id': user_id, 'provider': provider, 'token': self._encrypt_token(token), 'expires_at': token.get('expires_at') or int(time.time() + 3600), 'created_at': current_time, 'updated_at': current_time, } ) db.add(result) await db.commit() await db.refresh(result) if result: # Make a copy of the model data before closing session model = OAuthSessionModel( id=result.id, user_id=result.user_id, provider=result.provider, token=token, # Return decrypted token expires_at=result.expires_at, created_at=result.created_at, updated_at=result.updated_at, ) return model else: return None except Exception as e: log.error(f'Error creating OAuth session: {e}') return None async def get_session_by_id( self, session_id: str, db: Optional[AsyncSession] = None ) -> Optional[OAuthSessionModel]: """Get OAuth session by ID""" try: async with get_async_db_context(db) as db: result = await db.execute(select(OAuthSession).filter_by(id=session_id)) session = result.scalars().first() if session: return OAuthSessionModel( id=session.id, user_id=session.user_id, provider=session.provider, token=self._decrypt_token(session.token), expires_at=session.expires_at, created_at=session.created_at, updated_at=session.updated_at, ) return None except Exception as e: log.error(f'Error getting OAuth session by ID: {e}') return None async def get_session_by_id_and_user_id( self, session_id: str, user_id: str, db: Optional[AsyncSession] = None ) -> Optional[OAuthSessionModel]: """Get OAuth session by ID and user ID""" try: async with get_async_db_context(db) as db: result = await db.execute(select(OAuthSession).filter_by(id=session_id, user_id=user_id)) session = result.scalars().first() if session: return OAuthSessionModel( id=session.id, user_id=session.user_id, provider=session.provider, token=self._decrypt_token(session.token), expires_at=session.expires_at, created_at=session.created_at, updated_at=session.updated_at, ) return None except Exception as e: log.error(f'Error getting OAuth session by ID: {e}') return None async def get_session_by_provider_and_user_id( self, provider: str, user_id: str, db: Optional[AsyncSession] = None ) -> Optional[OAuthSessionModel]: """Get OAuth session by provider and user ID""" try: async with get_async_db_context(db) as db: result = await db.execute( select(OAuthSession) .filter_by(provider=provider, user_id=user_id) .order_by(OAuthSession.created_at.desc()) ) session = result.scalars().first() if session: return OAuthSessionModel( id=session.id, user_id=session.user_id, provider=session.provider, token=self._decrypt_token(session.token), expires_at=session.expires_at, created_at=session.created_at, updated_at=session.updated_at, ) return None except Exception as e: log.error(f'Error getting OAuth session by provider and user ID: {e}') return None async def get_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> List[OAuthSessionModel]: """Get all OAuth sessions for a user""" try: async with get_async_db_context(db) as db: result = await db.execute(select(OAuthSession).filter_by(user_id=user_id)) sessions = result.scalars().all() results = [] for session in sessions: try: results.append( OAuthSessionModel( id=session.id, user_id=session.user_id, provider=session.provider, token=self._decrypt_token(session.token), expires_at=session.expires_at, created_at=session.created_at, updated_at=session.updated_at, ) ) except Exception as e: log.warning( f'Skipping OAuth session {session.id} due to decryption failure, deleting corrupted session: {type(e).__name__}: {e}' ) await db.execute(delete(OAuthSession).filter_by(id=session.id)) await db.commit() return results except Exception as e: log.error(f'Error getting OAuth sessions by user ID: {e}') return [] async def update_session_by_id( self, session_id: str, token: dict, db: Optional[AsyncSession] = None ) -> Optional[OAuthSessionModel]: """Update OAuth session tokens""" try: async with get_async_db_context(db) as db: current_time = int(time.time()) await db.execute( update(OAuthSession) .filter_by(id=session_id) .values( token=self._encrypt_token(token), expires_at=token.get('expires_at') or int(time.time() + 3600), updated_at=current_time, ) ) await db.commit() result = await db.execute(select(OAuthSession).filter_by(id=session_id)) session = result.scalars().first() if session: return OAuthSessionModel( id=session.id, user_id=session.user_id, provider=session.provider, token=self._decrypt_token(session.token), expires_at=session.expires_at, created_at=session.created_at, updated_at=session.updated_at, ) return None except Exception as e: log.error(f'Error updating OAuth session tokens: {e}') return None async def delete_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> bool: """Delete an OAuth session""" try: async with get_async_db_context(db) as db: result = await db.execute(delete(OAuthSession).filter_by(id=session_id)) await db.commit() return result.rowcount > 0 except Exception as e: log.error(f'Error deleting OAuth session: {e}') return False async def delete_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool: """Delete all OAuth sessions for a user""" try: async with get_async_db_context(db) as db: await db.execute(delete(OAuthSession).filter_by(user_id=user_id)) await db.commit() return True except Exception as e: log.error(f'Error deleting OAuth sessions by user ID: {e}') return False async def delete_sessions_by_user_id_and_provider( self, user_id: str, provider: str, db: Optional[AsyncSession] = None ) -> bool: """Delete all OAuth sessions for a specific user and provider""" try: async with get_async_db_context(db) as db: result = await db.execute(delete(OAuthSession).filter_by(user_id=user_id, provider=provider)) await db.commit() return result.rowcount > 0 except Exception as e: log.error(f'Error deleting OAuth sessions for user {user_id} and provider {provider}: {e}') return False async def delete_sessions_by_provider(self, provider: str, db: Optional[AsyncSession] = None) -> bool: """Delete all OAuth sessions for a provider""" try: async with get_async_db_context(db) as db: await db.execute(delete(OAuthSession).filter_by(provider=provider)) await db.commit() return True except Exception as e: log.error(f'Error deleting OAuth sessions by provider {provider}: {e}') return False OAuthSessions = OAuthSessionTable()