r2rpy/py/core/providers/auth/r2r_auth.py

import logging
import os
from datetime import datetime, timedelta, timezone
from typing import Optional
from uuid import UUID

from fastapi import Depends, HTTPException
from fastapi.security import OAuth2PasswordBearer

from core.base import (
    AuthConfig,
    AuthProvider,
    CollectionResponse,
    CryptoProvider,
    EmailProvider,
    R2RException,
    Token,
    TokenData,
)
from core.base.api.models import User

from ..database import PostgresDatabaseProvider

DEFAULT_ACCESS_LIFETIME_IN_MINUTES = 3600
DEFAULT_REFRESH_LIFETIME_IN_DAYS = 7

logger = logging.getLogger()
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


def normalize_email(email: str) -> str:
    """Normalizes an email address by converting it to lowercase. This ensures
    consistent email handling throughout the application.

    Args:
        email: The email address to normalize

    Returns:
        The normalized (lowercase) email address
    """
    return email.lower() if email else ""


class R2RAuthProvider(AuthProvider):
    def __init__(
        self,
        config: AuthConfig,
        crypto_provider: CryptoProvider,
        database_provider: PostgresDatabaseProvider,
        email_provider: EmailProvider,
    ):
        super().__init__(
            config, crypto_provider, database_provider, email_provider
        )
        self.database_provider: PostgresDatabaseProvider = database_provider
        logger.debug(f"Initializing R2RAuthProvider with config: {config}")

        # We no longer use a local secret_key or defaults here.
        # All key handling is done in the crypto_provider.
        self.access_token_lifetime_in_minutes = (
            config.access_token_lifetime_in_minutes
            or os.getenv("R2R_ACCESS_LIFE_IN_MINUTES")
            or DEFAULT_ACCESS_LIFETIME_IN_MINUTES
        )
        self.refresh_token_lifetime_in_days = (
            config.refresh_token_lifetime_in_days
            or os.getenv("R2R_REFRESH_LIFE_IN_DAYS")
            or DEFAULT_REFRESH_LIFETIME_IN_DAYS
        )
        self.config: AuthConfig = config

    async def initialize(self):
        try:
            user = await self.register(
                email=normalize_email(self.admin_email),
                password=self.admin_password,
                is_superuser=True,
            )
            await self.database_provider.users_handler.mark_user_as_superuser(
                id=user.id
            )
        except R2RException:
            logger.info("Default admin user already exists.")

    def create_access_token(self, data: dict) -> str:
        expire = datetime.now(timezone.utc) + timedelta(
            minutes=float(self.access_token_lifetime_in_minutes)
        )
        # Add token_type and pass data/expiry to crypto_provider
        data_with_type = {**data, "token_type": "access"}
        return self.crypto_provider.generate_secure_token(
            data=data_with_type,
            expiry=expire,
        )

    def create_refresh_token(self, data: dict) -> str:
        expire = datetime.now(timezone.utc) + timedelta(
            days=float(self.refresh_token_lifetime_in_days)
        )
        data_with_type = {**data, "token_type": "refresh"}
        return self.crypto_provider.generate_secure_token(
            data=data_with_type,
            expiry=expire,
        )

    async def decode_token(self, token: str) -> TokenData:
        if "token=" in token:
            token = token.split("token=")[1]
        if "&tokenType=refresh" in token:
            token = token.split("&tokenType=refresh")[0]
        # First, check if the token is blacklisted
        if await self.database_provider.token_handler.is_token_blacklisted(
            token=token
        ):
            raise R2RException(
                status_code=401, message="Token has been invalidated"
            )

        # Verify token using crypto_provider
        payload = self.crypto_provider.verify_secure_token(token=token)
        if payload is None:
            raise R2RException(
                status_code=401, message="Invalid or expired token"
            )

        email = payload.get("sub")
        token_type = payload.get("token_type")
        exp = payload.get("exp")

        if email is None or token_type is None or exp is None:
            raise R2RException(status_code=401, message="Invalid token claims")

        email_str: str = email
        token_type_str: str = token_type
        exp_float: float = exp

        exp_datetime = datetime.fromtimestamp(exp_float, tz=timezone.utc)
        if exp_datetime < datetime.now(timezone.utc):
            raise R2RException(status_code=401, message="Token has expired")

        return TokenData(
            email=normalize_email(email_str),
            token_type=token_type_str,
            exp=exp_datetime,
        )

    async def authenticate_api_key(self, api_key: str) -> User:
        """Authenticate using an API key of the form "public_key.raw_key".

        Returns a User if successful, or raises R2RException if not.
        """
        try:
            key_id, raw_key = api_key.split(".", 1)
        except ValueError as e:
            raise R2RException(
                status_code=401, message="Invalid API key format"
            ) from e

        key_record = (
            await self.database_provider.users_handler.get_api_key_record(
                key_id=key_id
            )
        )
        if not key_record:
            raise R2RException(status_code=401, message="Invalid API key")

        if not self.crypto_provider.verify_api_key(
            raw_api_key=raw_key, hashed_key=key_record["hashed_key"]
        ):
            raise R2RException(status_code=401, message="Invalid API key")

        user = await self.database_provider.users_handler.get_user_by_id(
            id=key_record["user_id"]
        )
        if not user.is_active:
            raise R2RException(
                status_code=401, message="User account is inactive"
            )

        return user

    async def user(self, token: str = Depends(oauth2_scheme)) -> User:
        """Attempt to authenticate via JWT first, then fallback to API key."""
        # Try JWT auth
        try:
            token_data = await self.decode_token(token=token)
            if not token_data.email:
                raise R2RException(
                    status_code=401, message="Could not validate credentials"
                )
            user = (
                await self.database_provider.users_handler.get_user_by_email(
                    email=normalize_email(token_data.email)
                )
            )
            if user is None:
                raise R2RException(
                    status_code=401,
                    message="Invalid authentication credentials",
                )
            return user
        except R2RException:
            # If JWT fails, try API key auth
            # OAuth2PasswordBearer provides token as "Bearer xxx", strip it if needed
            token = token.removeprefix("Bearer ")
            return await self.authenticate_api_key(api_key=token)

    def get_current_active_user(
        self, current_user: User = Depends(user)
    ) -> User:
        if not current_user.is_active:
            raise R2RException(status_code=400, message="Inactive user")
        return current_user

    async def register(
        self,
        email: str,
        password: Optional[str] = None,
        is_superuser: bool = False,
        is_verified: bool = False,
        account_type: str = "password",
        github_id: Optional[str] = None,
        google_id: Optional[str] = None,
        name: Optional[str] = None,
        bio: Optional[str] = None,
        profile_picture: Optional[str] = None,
    ) -> User:
        if account_type == "password":
            if not password:
                raise R2RException(
                    status_code=400,
                    message="Password is required for password accounts",
                )
        else:
            if github_id and google_id:
                raise R2RException(
                    status_code=400,
                    message="Cannot register OAuth with both GitHub and Google IDs",
                )
            if not github_id and not google_id:
                raise R2RException(
                    status_code=400,
                    message="Invalid OAuth specification without GitHub or Google ID",
                )
        new_user = await self.database_provider.users_handler.create_user(
            email=normalize_email(email),
            password=password,
            is_superuser=is_superuser,
            is_verified=is_verified,
            account_type=account_type,
            github_id=github_id,
            google_id=google_id,
            name=name,
            bio=bio,
            profile_picture=profile_picture,
        )
        default_collection: CollectionResponse = (
            await self.database_provider.collections_handler.create_collection(
                owner_id=new_user.id,
            )
        )
        await self.database_provider.graphs_handler.create(
            collection_id=default_collection.id,
            name=default_collection.name,
            description=default_collection.description,
        )

        await self.database_provider.users_handler.add_user_to_collection(
            new_user.id, default_collection.id
        )

        new_user = await self.database_provider.users_handler.get_user_by_id(
            new_user.id
        )

        if self.config.require_email_verification and not is_verified:
            verification_code, _ = await self.send_verification_email(
                email=normalize_email(email), user=new_user
            )

        return new_user

    async def send_verification_email(
        self, email: str, user: Optional[User] = None
    ) -> tuple[str, datetime]:
        if user is None:
            user = (
                await self.database_provider.users_handler.get_user_by_email(
                    email=normalize_email(email)
                )
            )
            if not user:
                raise R2RException(status_code=404, message="User not found")

        verification_code = self.crypto_provider.generate_verification_code()
        expiry = datetime.now(timezone.utc) + timedelta(hours=24)

        await self.database_provider.users_handler.store_verification_code(
            id=user.id,
            verification_code=verification_code,
            expiry=expiry,
        )

        if hasattr(user, "verification_code_expiry"):
            user.verification_code_expiry = expiry

        first_name = (
            user.name.split(" ")[0] if user.name else email.split("@")[0]
        )

        await self.email_provider.send_verification_email(
            to_email=user.email,
            verification_code=verification_code,
            dynamic_template_data={"first_name": first_name},
        )

        return verification_code, expiry

    async def verify_email(
        self, email: str, verification_code: str
    ) -> dict[str, str]:
        user_id = await self.database_provider.users_handler.get_user_id_by_verification_code(
            verification_code=verification_code
        )
        await self.database_provider.users_handler.mark_user_as_verified(
            id=user_id
        )
        await self.database_provider.users_handler.remove_verification_code(
            verification_code=verification_code
        )
        return {"message": "Email verified successfully"}

    async def login(self, email: str, password: str) -> dict[str, Token]:
        logger.debug(f"Attempting login for email: {email}")
        user = await self.database_provider.users_handler.get_user_by_email(
            email=normalize_email(email)
        )

        if user.account_type != "password":
            logger.warning(
                f"Password login not allowed for {user.account_type} accounts: {email}"
            )
            raise R2RException(
                status_code=401,
                message=f"This account is configured for {user.account_type} login, not password.",
            )

        logger.debug(f"User found: {user}")

        if not isinstance(user.hashed_password, str):
            logger.error(
                f"Invalid hashed_password type: {type(user.hashed_password)}"
            )
            raise HTTPException(
                status_code=500,
                detail="Invalid password hash in database",
            )

        try:
            password_verified = self.crypto_provider.verify_password(
                plain_password=password,
                hashed_password=user.hashed_password,
            )
        except Exception as e:
            logger.error(f"Error during password verification: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail="Error during password verification",
            ) from e

        if not password_verified:
            logger.warning(f"Invalid password for user: {email}")
            raise R2RException(
                status_code=401, message="Incorrect email or password"
            )

        if not user.is_verified and self.config.require_email_verification:
            logger.warning(f"Unverified user attempted login: {email}")
            raise R2RException(status_code=401, message="Email not verified")

        access_token = self.create_access_token(
            data={"sub": normalize_email(user.email)}
        )
        refresh_token = self.create_refresh_token(
            data={"sub": normalize_email(user.email)}
        )
        return {
            "access_token": Token(token=access_token, token_type="access"),
            "refresh_token": Token(token=refresh_token, token_type="refresh"),
        }

    async def refresh_access_token(
        self, refresh_token: str
    ) -> dict[str, Token]:
        token_data = await self.decode_token(refresh_token)
        if token_data.token_type != "refresh":
            raise R2RException(
                status_code=401, message="Invalid refresh token"
            )

        # Invalidate the old refresh token and create a new one
        await self.database_provider.token_handler.blacklist_token(
            token=refresh_token
        )

        new_access_token = self.create_access_token(
            data={"sub": normalize_email(token_data.email)}
        )
        new_refresh_token = self.create_refresh_token(
            data={"sub": normalize_email(token_data.email)}
        )
        return {
            "access_token": Token(token=new_access_token, token_type="access"),
            "refresh_token": Token(
                token=new_refresh_token, token_type="refresh"
            ),
        }

    async def change_password(
        self, user: User, current_password: str, new_password: str
    ) -> dict[str, str]:
        if not isinstance(user.hashed_password, str):
            logger.error(
                f"Invalid hashed_password type: {type(user.hashed_password)}"
            )
            raise HTTPException(
                status_code=500,
                detail="Invalid password hash in database",
            )

        if not self.crypto_provider.verify_password(
            plain_password=current_password,
            hashed_password=user.hashed_password,
        ):
            raise R2RException(
                status_code=400, message="Incorrect current password"
            )

        hashed_new_password = self.crypto_provider.get_password_hash(
            password=new_password
        )
        await self.database_provider.users_handler.update_user_password(
            id=user.id,
            new_hashed_password=hashed_new_password,
        )
        try:
            await self.email_provider.send_password_changed_email(
                to_email=normalize_email(user.email),
                dynamic_template_data={
                    "first_name": (
                        user.name.split(" ")[0] or "User"
                        if user.name
                        else "User"
                    )
                },
            )
        except Exception as e:
            logger.error(
                f"Failed to send password change notification: {str(e)}"
            )

        return {"message": "Password changed successfully"}

    async def request_password_reset(self, email: str) -> dict[str, str]:
        try:
            user = (
                await self.database_provider.users_handler.get_user_by_email(
                    email=normalize_email(email)
                )
            )

            reset_token = self.crypto_provider.generate_verification_code()
            expiry = datetime.now(timezone.utc) + timedelta(hours=1)
            await self.database_provider.users_handler.store_reset_token(
                id=user.id,
                reset_token=reset_token,
                expiry=expiry,
            )

            first_name = (
                user.name.split(" ")[0] if user.name else email.split("@")[0]
            )
            await self.email_provider.send_password_reset_email(
                to_email=normalize_email(email),
                reset_token=reset_token,
                dynamic_template_data={"first_name": first_name},
            )

            return {
                "message": "If the email exists, a reset link has been sent"
            }
        except R2RException as e:
            if e.status_code == 404:
                # User doesn't exist; return a success message anyway
                return {
                    "message": "If the email exists, a reset link has been sent"
                }
            else:
                raise

    async def confirm_password_reset(
        self, reset_token: str, new_password: str
    ) -> dict[str, str]:
        user_id = await self.database_provider.users_handler.get_user_id_by_reset_token(
            reset_token=reset_token
        )
        if not user_id:
            raise R2RException(
                status_code=400, message="Invalid or expired reset token"
            )

        hashed_new_password = self.crypto_provider.get_password_hash(
            password=new_password
        )
        await self.database_provider.users_handler.update_user_password(
            id=user_id,
            new_hashed_password=hashed_new_password,
        )
        await self.database_provider.users_handler.remove_reset_token(
            id=user_id
        )
        # Get the user information
        user = await self.database_provider.users_handler.get_user_by_id(
            id=user_id
        )

        try:
            await self.email_provider.send_password_changed_email(
                to_email=normalize_email(user.email),
                dynamic_template_data={
                    "first_name": (
                        user.name.split(" ")[0] or "User"
                        if user.name
                        else "User"
                    )
                },
            )
        except Exception as e:
            logger.error(
                f"Failed to send password change notification: {str(e)}"
            )

        return {"message": "Password reset successfully"}

    async def logout(self, token: str) -> dict[str, str]:
        await self.database_provider.token_handler.blacklist_token(token=token)
        return {"message": "Logged out successfully"}

    async def clean_expired_blacklisted_tokens(self):
        await self.database_provider.token_handler.clean_expired_blacklisted_tokens()

    async def send_reset_email(self, email: str) -> dict:
        verification_code, expiry = await self.send_verification_email(
            email=normalize_email(email)
        )

        return {
            "verification_code": verification_code,
            "expiry": expiry,
            "message": f"Verification email sent successfully to {email}",
        }

    async def create_user_api_key(
        self,
        user_id: UUID,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> dict[str, str]:
        key_id, raw_api_key = self.crypto_provider.generate_api_key()
        hashed_key = self.crypto_provider.hash_api_key(raw_api_key)

        api_key_uuid = (
            await self.database_provider.users_handler.store_user_api_key(
                user_id=user_id,
                key_id=key_id,
                hashed_key=hashed_key,
                name=name,
                description=description,
            )
        )

        return {
            "api_key": f"{key_id}.{raw_api_key}",
            "key_id": str(api_key_uuid),
            "public_key": key_id,
            "name": name or "",
        }

    async def list_user_api_keys(self, user_id: UUID) -> list[dict]:
        return await self.database_provider.users_handler.get_user_api_keys(
            user_id=user_id
        )

    async def delete_user_api_key(self, user_id: UUID, key_id: UUID) -> bool:
        return await self.database_provider.users_handler.delete_api_key(
            user_id=user_id,
            key_id=key_id,
        )

    async def rename_api_key(
        self, user_id: UUID, key_id: UUID, new_name: str
    ) -> bool:
        return await self.database_provider.users_handler.update_api_key_name(
            user_id=user_id,
            key_id=key_id,
            name=new_name,
        )

    async def oauth_callback_handler(
        self, provider: str, oauth_id: str, email: str
    ) -> dict[str, Token]:
        """Handles a login/registration flow for OAuth providers (e.g., Google
        or GitHub).

        :param provider: "google" or "github"
        :param oauth_id: The unique ID from the OAuth provider (e.g. Google's
            'sub')
        :param email: The user's email from the provider, if available.
        :return: dict with access_token and refresh_token
        """
        # 1) Attempt to find user by google_id or github_id, or by email
        #    The logic depends on your preference. We'll assume "google" => google_id, etc.
        try:
            if provider == "google":
                try:
                    user = await self.database_provider.users_handler.get_user_by_email(
                        normalize_email(email)
                    )
                    # If user found, check if user.google_id matches or is null. If null, update it
                    if user and not user.google_id:
                        raise R2RException(
                            status_code=401,
                            message="User already exists and is not linked to Google account",
                        )
                except Exception:
                    # Create new user
                    user = await self.register(
                        email=normalize_email(email)
                        or f"{oauth_id}@google_oauth.fake",  # fallback
                        password=None,  # no password
                        account_type="oauth",
                        google_id=oauth_id,
                    )
            elif provider == "github":
                try:
                    user = await self.database_provider.users_handler.get_user_by_email(
                        normalize_email(email)
                    )
                    # If user found, check if user.google_id matches or is null. If null, update it
                    if user and not user.github_id:
                        raise R2RException(
                            status_code=401,
                            message="User already exists and is not linked to Github account",
                        )
                except Exception:
                    # Create new user
                    user = await self.register(
                        email=normalize_email(email)
                        or f"{oauth_id}@github_oauth.fake",  # fallback
                        password=None,  # no password
                        account_type="oauth",
                        github_id=oauth_id,
                    )
            # else handle other providers

        except R2RException:
            # If no user found or creation fails
            raise R2RException(
                status_code=401, message="Could not create or fetch user"
            ) from None

        # If user is inactive, etc.
        if not user.is_active:
            raise R2RException(
                status_code=401, message="User account is inactive"
            )

        # Possibly mark user as verified if you trust the OAuth provider's email
        user.is_verified = True
        await self.database_provider.users_handler.update_user(user)

        # 2) Generate tokens
        access_token = self.create_access_token(
            data={"sub": normalize_email(user.email)}
        )
        refresh_token = self.create_refresh_token(
            data={"sub": normalize_email(user.email)}
        )

        return {
            "access_token": Token(token=access_token, token_type="access"),
            "refresh_token": Token(token=refresh_token, token_type="refresh"),
        }