auth.py

"""Define the Autorization Manager."""

import datetime

import jwt
from fastapi import BackgroundTasks, Depends, HTTPException, Request, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import NameEmail
from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession

from app.config.settings import get_settings
from app.database.db import get_database
from app.database.helpers import (
    get_user_by_email_,
    get_user_by_id_,
    hash_password,
)
from app.managers.email import EmailManager
from app.managers.helpers import MAX_JWT_TOKEN_LENGTH, is_valid_jwt_format
from app.models.enums import RoleType
from app.models.user import User
from app.schemas.email import EmailTemplateSchema
from app.schemas.request.auth import TokenRefreshRequest


class ResponseMessages:
    """Error strings for different circumstances."""

    CANT_GENERATE_JWT = "Unable to generate the JWT"
    CANT_GENERATE_REFRESH = "Unable to generate the Refresh Token"
    CANT_GENERATE_VERIFY = "Unable to generate the Verification Token"
    CANT_GENERATE_RESET = "Unable to generate the Password Reset Token"
    INVALID_TOKEN = "That token is Invalid"  # noqa: S105
    EXPIRED_TOKEN = "That token has Expired"  # noqa: S105
    VERIFICATION_SUCCESS = "User successfully Verified"
    USER_NOT_FOUND = "User not Found"
    ALREADY_VALIDATED = "You are already validated"
    VALIDATION_RESENT = "Validation email re-sent"
    RESET_EMAIL_SENT = "Password reset email sent if user exists"
    PASSWORD_RESET_SUCCESS = "Password successfully reset"  # noqa: S105


class AuthManager:
    """Handle the JWT Auth."""

    @staticmethod
    def encode_token(user: User) -> str:
        """Create and return a JTW token."""
        try:
            payload = {
                "sub": user.id,
                "exp": datetime.datetime.now(tz=datetime.timezone.utc)
                + datetime.timedelta(
                    minutes=get_settings().access_token_expire_minutes
                ),
            }
            return jwt.encode(
                payload, get_settings().secret_key, algorithm="HS256"
            )
        except (jwt.PyJWTError, AttributeError) as exc:
            # log the exception
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.CANT_GENERATE_JWT
            ) from exc

    @staticmethod
    def encode_refresh_token(user: User) -> str:
        """Create and return a JTW token."""
        try:
            payload = {
                "sub": user.id,
                "exp": datetime.datetime.now(tz=datetime.timezone.utc)
                + datetime.timedelta(minutes=60 * 24 * 30),
                "typ": "refresh",
            }
            return jwt.encode(
                payload, get_settings().secret_key, algorithm="HS256"
            )
        except (jwt.PyJWTError, AttributeError) as exc:
            # log the exception
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED,
                ResponseMessages.CANT_GENERATE_REFRESH,
            ) from exc

    @staticmethod
    def encode_verify_token(user: User) -> str:
        """Create and return a JTW token."""
        try:
            payload = {
                "sub": user.id,
                "exp": datetime.datetime.now(tz=datetime.timezone.utc)
                + datetime.timedelta(minutes=10),
                "typ": "verify",
            }
            return jwt.encode(
                payload, get_settings().secret_key, algorithm="HS256"
            )
        except (jwt.PyJWTError, AttributeError) as exc:
            # log the exception
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED,
                ResponseMessages.CANT_GENERATE_VERIFY,
            ) from exc

    @staticmethod
    def encode_reset_token(user: User) -> str:
        """Create and return a password reset JWT token."""
        try:
            payload = {
                "sub": user.id,
                "exp": datetime.datetime.now(tz=datetime.timezone.utc)
                + datetime.timedelta(minutes=30),
                "typ": "reset",
            }
            return jwt.encode(
                payload, get_settings().secret_key, algorithm="HS256"
            )
        except (jwt.PyJWTError, AttributeError) as exc:
            # log the exception
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED,
                ResponseMessages.CANT_GENERATE_RESET,
            ) from exc

    @staticmethod
    async def refresh(
        refresh_token: TokenRefreshRequest, session: AsyncSession
    ) -> str:
        """Refresh an expired JWT token, given a valid Refresh token."""
        # Validate token format before processing
        if (
            not refresh_token.refresh
            or len(refresh_token.refresh) > MAX_JWT_TOKEN_LENGTH
            or not is_valid_jwt_format(refresh_token.refresh)
        ):
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            )

        try:
            payload = jwt.decode(
                refresh_token.refresh,
                get_settings().secret_key,
                algorithms=["HS256"],
                options={"verify_sub": False},
            )

            if payload["typ"] != "refresh":
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            user_data = await get_user_by_id_(payload["sub"], session)

            if not user_data:
                raise HTTPException(
                    status.HTTP_404_NOT_FOUND, ResponseMessages.USER_NOT_FOUND
                )

            # block a banned user
            if bool(user_data.banned):
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )
            new_token = AuthManager.encode_token(user_data)

        except jwt.ExpiredSignatureError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.EXPIRED_TOKEN
            ) from exc
        except jwt.InvalidTokenError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            ) from exc
        else:
            return new_token

    @staticmethod
    async def verify(code: str, session: AsyncSession) -> None:
        """Verify a new User's Email using the token they were sent."""
        # Validate token format before processing
        if (
            not code
            or len(code) > MAX_JWT_TOKEN_LENGTH
            or not is_valid_jwt_format(code)
        ):
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            )

        try:
            payload = jwt.decode(
                code,
                get_settings().secret_key,
                algorithms=["HS256"],
                options={"verify_sub": False},
            )

            user_data = await session.get(User, payload["sub"])

            if not user_data:
                raise HTTPException(
                    status.HTTP_404_NOT_FOUND, ResponseMessages.USER_NOT_FOUND
                )

            if payload["typ"] != "verify":
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            # block a banned user
            if bool(user_data.banned):
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            if bool(user_data.verified):
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            await session.execute(
                update(User)
                .where(User.id == payload["sub"])
                .values(
                    verified=True,
                )
            )
            await session.commit()

            raise HTTPException(
                status.HTTP_200_OK, ResponseMessages.VERIFICATION_SUCCESS
            )

        except jwt.ExpiredSignatureError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.EXPIRED_TOKEN
            ) from exc
        except jwt.InvalidTokenError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            ) from exc

    @staticmethod
    async def forgot_password(
        email: str, background_tasks: BackgroundTasks, session: AsyncSession
    ) -> None:
        """Send a password reset email to the user if they exist."""
        # Get user by email - but don't reveal if user exists for security
        user = await get_user_by_email_(email, session)

        # Always return success message to prevent email enumeration
        if not user:
            return

        # Don't send reset email to banned users
        if bool(user.banned):
            return

        # Generate reset token
        reset_token = AuthManager.encode_reset_token(user)

        # Send password reset email
        email_manager = EmailManager()
        email_manager.template_send(
            background_tasks,
            EmailTemplateSchema(
                recipients=[NameEmail(name=user.first_name, email=user.email)],
                subject=f"{get_settings().api_title} - Password Reset",
                body={
                    "name": user.first_name,
                    "application": get_settings().api_title,
                    "base_url": get_settings().base_url,
                    "reset_token": reset_token,
                },
                template_name="password_reset.html",
            ),
        )

    @staticmethod
    async def reset_password(
        code: str, new_password: str, session: AsyncSession
    ) -> None:
        """Reset a user's password using the reset token."""
        try:
            payload = jwt.decode(
                code,
                get_settings().secret_key,
                algorithms=["HS256"],
                options={"verify_sub": False},
            )

            user_data = await session.get(User, payload["sub"])

            if not user_data:
                raise HTTPException(
                    status.HTTP_404_NOT_FOUND, ResponseMessages.USER_NOT_FOUND
                )

            if payload["typ"] != "reset":
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            # Block banned users from resetting password
            if bool(user_data.banned):
                raise HTTPException(
                    status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
                )

            # Hash the new password
            hashed_password = hash_password(new_password)

            # Update the user's password
            await session.execute(
                update(User)
                .where(User.id == payload["sub"])
                .values(password=hashed_password)
            )
            await session.commit()

        except jwt.ExpiredSignatureError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.EXPIRED_TOKEN
            ) from exc
        except jwt.InvalidTokenError as exc:
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            ) from exc

    @staticmethod
    async def resend_verify_code(
        user: int, background_tasks: BackgroundTasks, session: AsyncSession
    ) -> None:  # pragma: no cover (code not used at this time)
        """Resend the user a verification email."""
        user_data = await get_user_by_id_(user, session)

        if not user_data:
            raise HTTPException(
                status.HTTP_404_NOT_FOUND, ResponseMessages.USER_NOT_FOUND
            )

        # block a banned user
        if bool(user_data.banned):
            raise HTTPException(
                status.HTTP_401_UNAUTHORIZED, ResponseMessages.INVALID_TOKEN
            )

        if bool(user_data.verified):
            raise HTTPException(
                status.HTTP_400_BAD_REQUEST,
                ResponseMessages.ALREADY_VALIDATED,
            )

        email = EmailManager()
        user_full_name = f"{user_data.first_name} {user_data.last_name}"
        email.template_send(
            background_tasks,
            EmailTemplateSchema(
                recipients=[NameEmail(user_full_name, user_data.email)],
                subject=f"Welcome to {get_settings().api_title}!",
                body={
                    "application": f"{get_settings().api_title}",
                    "user": user_data.email,
                    "base_url": get_settings().base_url,
                    "verification": AuthManager.encode_verify_token(user_data),
                },
                template_name="welcome.html",
            ),
        )
        # await email.simple_send(
        #     EmailSchema(
        #         recipients=[user_data["email"]],
        #         subject=f"Welcome to {get_settings().api_title}!",
        #         body="Test Email",
        #     ),
        # )

        raise HTTPException(
            status.HTTP_200_OK, ResponseMessages.VALIDATION_RESENT
        )


bearer = HTTPBearer(auto_error=False)


async def get_jwt_user(
    request: Request,
    db: AsyncSession = Depends(get_database),
    credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
) -> User | None:
    """Get user from JWT token."""
    if not credentials:
        return None

    try:
        # Decode and validate the token
        payload = jwt.decode(
            credentials.credentials,
            get_settings().secret_key,
            algorithms=["HS256"],
            options={"verify_sub": False},
        )
        user_data = await get_user_by_id_(payload["sub"], db)

        # Check user validity - user must exist, be verified, and not banned
        if not user_data:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ResponseMessages.INVALID_TOKEN,
            )

        if bool(user_data.banned) or not bool(user_data.verified):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ResponseMessages.INVALID_TOKEN,
            )

        # Store user in request state
        request.state.user = user_data

    except jwt.ExpiredSignatureError as exc:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ResponseMessages.EXPIRED_TOKEN,
        ) from exc
    except jwt.InvalidTokenError as exc:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ResponseMessages.INVALID_TOKEN,
        ) from exc
    else:
        return user_data


oauth2_schema = get_jwt_user


def is_admin(request: Request) -> None:
    """Block if user is not an Admin."""
    if request.state.user.role != RoleType.admin:
        raise HTTPException(status.HTTP_403_FORBIDDEN, "Forbidden")


def can_edit_user(request: Request) -> None:
    """Check if the user can edit this resource.

    True if they own the resource or are Admin
    """
    if (
        request.state.user.role != RoleType.admin
        and request.state.user.id != int(request.path_params["user_id"])
    ):
        raise HTTPException(status.HTTP_403_FORBIDDEN, "Forbidden")


def is_banned(request: Request) -> None:
    """Dont let banned users access the route."""
    if request.state.user.banned:
        raise HTTPException(status.HTTP_403_FORBIDDEN, "Banned!")