blog.pyaqa.ru/app/presentation/web/flash.py

"""Flash messages middleware for web UI.

This module provides flash message functionality for the web interface,
allowing temporary messages to be passed between requests (e.g., after redirect).
Uses signed cookies for security.
"""

from fastapi import Request, Response
from itsdangerous import URLSafeSerializer

from app.infrastructure.config.settings import settings

FLASH_COOKIE_NAME = "flash_messages"
_SECRET_KEY = (
    settings.security.secret_key.get_secret_value()
    if hasattr(settings.security.secret_key, "get_secret_value")
    else settings.security.secret_key
)
SERIALIZER = URLSafeSerializer(_SECRET_KEY)


class FlashMessage:
    """Flash message model.

    Represents a single flash message with type and content.

    Attributes:
        message: The message text.
        category: Message category (success, error, warning, info).
    """

    def __init__(self, message: str, category: str = "info") -> None:
        """Initialize flash message.

        Args:
            message: The message text.
            category: Message category (success, error, warning, info).
        """
        self.message = message
        self.category = category

    def to_dict(self) -> dict[str, str]:
        """Convert to dictionary.

        Returns:
            Dictionary with message and category.
        """
        return {"message": self.message, "category": self.category}


class FlashManager:
    """Manager for flash messages.

    Handles storing and retrieving flash messages from cookies.
    Messages are cleared after being read.

    Attributes:
        request: FastAPI request object.
        messages: List of current messages.
    """

    CATEGORIES = {
        "success": "success",
        "error": "error",
        "warning": "warning",
        "info": "info",
    }

    def __init__(self, request: Request) -> None:
        """Initialize flash manager.

        Args:
            request: FastAPI request object.
        """
        self.request = request
        self.messages: list[FlashMessage] = []
        self._load_messages()

    def _load_messages(self) -> None:
        """Load messages from cookie."""
        cookie_value = self.request.cookies.get(FLASH_COOKIE_NAME)
        if cookie_value:
            try:
                data = SERIALIZER.loads(cookie_value)
                if isinstance(data, list):
                    self.messages = [FlashMessage(msg["message"], msg["category"]) for msg in data]
            except Exception:
                self.messages = []

    def add(self, message: str, category: str = "info") -> None:
        """Add a flash message.

        Args:
            message: The message text.
            category: Message category (success, error, warning, info).
        """
        self.messages.append(FlashMessage(message, category))

    def get_messages(self) -> list[dict[str, str]]:
        """Get all messages and clear them.

        Returns:
            List of message dictionaries.
        """
        result = [msg.to_dict() for msg in self.messages]
        self.messages = []
        return result

    def has_messages(self) -> bool:
        """Check if there are any messages.

        Returns:
            True if there are messages.
        """
        return len(self.messages) > 0

    def set_cookie(self, response: Response) -> None:
        """Set flash cookie on response.

        Args:
            response: FastAPI response object.
        """
        if self.messages:
            data = [msg.to_dict() for msg in self.messages]
            cookie_value = SERIALIZER.dumps(data)
            response.set_cookie(
                key=FLASH_COOKIE_NAME,
                value=cookie_value,
                httponly=True,
                secure=not settings.is_dev,
                samesite="lax",
                max_age=300,  # 5 minutes
            )
        else:
            response.delete_cookie(key=FLASH_COOKIE_NAME)


def flash(request: Request, message: str, category: str = "info") -> None:
    """Add flash message to request state.

    Convenience function to add flash message.
    Must be called before response is created.

    Args:
        request: FastAPI request object.
        message: The message text.
        category: Message category.
    """
    if not hasattr(request.state, "flash_manager"):
        request.state.flash_manager = FlashManager(request)
    request.state.flash_manager.add(message, category)


def get_flash_messages(request: Request) -> list[dict[str, str]]:
    """Get flash messages from request.

    Args:
        request: FastAPI request object.

    Returns:
        List of flash message dictionaries.
    """
    if hasattr(request.state, "flash_manager"):
        return request.state.flash_manager.get_messages()  # type: ignore[no-any-return]
    return []


async def setup_flash_manager(request: Request) -> None:
    """Setup flash manager on request state.

    Args:
        request: FastAPI request object.
    """
    if not hasattr(request.state, "flash_manager"):
        request.state.flash_manager = FlashManager(request)