core/auth/tokenstorage.py

from datetime import datetime, timedelta, timezone
import json
import time
from typing import Dict, Any, Optional, Tuple, List

from auth.jwtcodec import JWTCodec
from auth.validations import AuthInput
from services.redis import redis
from settings import ONETIME_TOKEN_LIFE_SPAN, SESSION_TOKEN_LIFE_SPAN
from utils.logger import root_logger as logger


class TokenStorage:
    """
    Класс для работы с хранилищем токенов в Redis
    """

    @staticmethod
    def _make_token_key(user_id: str, username: str, token: str) -> str:
        """
        Создает ключ для хранения токена

        Args:
            user_id: ID пользователя
            username: Имя пользователя
            token: Токен

        Returns:
            str: Ключ токена
        """
        # Сохраняем в старом формате для обратной совместимости
        return f"{user_id}-{username}-{token}"

    @staticmethod
    def _make_session_key(user_id: str, token: str) -> str:
        """
        Создает ключ в новом формате SessionManager

        Args:
            user_id: ID пользователя
            token: Токен

        Returns:
            str: Ключ сессии
        """
        return f"session:{user_id}:{token}"

    @staticmethod
    def _make_user_sessions_key(user_id: str) -> str:
        """
        Создает ключ для списка сессий пользователя

        Args:
            user_id: ID пользователя

        Returns:
            str: Ключ списка сессий
        """
        return f"user_sessions:{user_id}"

    @classmethod
    async def create_session(cls, user_id: str, username: str, device_info: Optional[Dict[str, str]] = None) -> str:
        """
        Создает новую сессию для пользователя

        Args:
            user_id: ID пользователя
            username: Имя пользователя
            device_info: Информация об устройстве (опционально)

        Returns:
            str: Токен сессии
        """
        logger.debug(f"[TokenStorage.create_session] Начало создания сессии для пользователя {user_id}")

        # Генерируем JWT токен с явным указанием времени истечения
        expiration_date = datetime.now(tz=timezone.utc) + timedelta(days=30)
        token = JWTCodec.encode({"id": user_id, "email": username}, exp=expiration_date)
        logger.debug(f"[TokenStorage.create_session] Создан JWT токен длиной {len(token)}")

        # Формируем ключи для Redis
        token_key = cls._make_token_key(user_id, username, token)
        logger.debug(f"[TokenStorage.create_session] Сформированы ключи: token_key={token_key}")
        
        # Формируем ключи в новом формате SessionManager для совместимости
        session_key = cls._make_session_key(user_id, token)
        user_sessions_key = cls._make_user_sessions_key(user_id)

        # Готовим данные для сохранения
        token_data = {
            "user_id": user_id,
            "username": username,
            "created_at": time.time(),
            "expires_at": time.time() + 30 * 24 * 60 * 60  # 30 дней
        }
        
        if device_info:
            token_data.update(device_info)
            
        logger.debug(f"[TokenStorage.create_session] Сформированы данные сессии: {token_data}")

        # Сохраняем в Redis старый формат
        pipeline = redis.pipeline()
        pipeline.hset(token_key, mapping=token_data)
        pipeline.expire(token_key, 30 * 24 * 60 * 60)  # 30 дней
        
        # Также сохраняем в новом формате SessionManager для обеспечения совместимости
        pipeline.hset(session_key, mapping=token_data)
        pipeline.expire(session_key, 30 * 24 * 60 * 60)  # 30 дней
        pipeline.sadd(user_sessions_key, token)
        pipeline.expire(user_sessions_key, 30 * 24 * 60 * 60)  # 30 дней
        
        results = await pipeline.execute()
        logger.info(f"[TokenStorage.create_session] Сессия успешно создана для пользователя {user_id}")

        return token

    @classmethod
    async def exists(cls, token_key: str) -> bool:
        """
        Проверяет существование токена по ключу

        Args:
            token_key: Ключ токена

        Returns:
            bool: True, если токен существует
        """
        exists = await redis.exists(token_key)
        return bool(exists)

    @classmethod
    async def validate_token(cls, token: str) -> Tuple[bool, Optional[Dict[str, Any]]]:
        """
        Проверяет валидность токена

        Args:
            token: JWT токен

        Returns:
            Tuple[bool, Dict[str, Any]]: (Валиден ли токен, данные токена)
        """
        try:
            # Декодируем JWT токен
            payload = JWTCodec.decode(token)
            if not payload:
                logger.warning(f"[TokenStorage.validate_token] Токен не валиден (не удалось декодировать)")
                return False, None
                
            user_id = payload.user_id
            username = payload.username
            
            # Формируем ключи для Redis в обоих форматах
            token_key = cls._make_token_key(user_id, username, token)
            session_key = cls._make_session_key(user_id, token)
            
            # Проверяем в обоих форматах для совместимости
            old_exists = await redis.exists(token_key)
            new_exists = await redis.exists(session_key)
            
            if old_exists or new_exists:
                logger.info(f"[TokenStorage.validate_token] Токен валиден для пользователя {user_id}")
                
                # Получаем данные токена из актуального хранилища
                if new_exists:
                    token_data = await redis.hgetall(session_key)
                else:
                    token_data = await redis.hgetall(token_key)
                    
                    # Если найден только в старом формате, создаем запись в новом формате
                    if not new_exists:
                        logger.info(f"[TokenStorage.validate_token] Миграция токена в новый формат: {session_key}")
                        await redis.hset(session_key, mapping=token_data)
                        await redis.expire(session_key, 30 * 24 * 60 * 60)
                        await redis.sadd(cls._make_user_sessions_key(user_id), token)
                
                return True, token_data
            else:
                logger.warning(f"[TokenStorage.validate_token] Токен не найден в Redis: {token_key}")
                return False, None
                
        except Exception as e:
            logger.error(f"[TokenStorage.validate_token] Ошибка при проверке токена: {e}")
            return False, None

    @classmethod
    async def invalidate_token(cls, token: str) -> bool:
        """
        Инвалидирует токен

        Args:
            token: JWT токен

        Returns:
            bool: True, если токен успешно инвалидирован
        """
        try:
            # Декодируем JWT токен
            payload = JWTCodec.decode(token)
            if not payload:
                logger.warning(f"[TokenStorage.invalidate_token] Токен не валиден (не удалось декодировать)")
                return False
                
            user_id = payload.user_id
            username = payload.username
            
            # Формируем ключи для Redis в обоих форматах
            token_key = cls._make_token_key(user_id, username, token)
            session_key = cls._make_session_key(user_id, token)
            user_sessions_key = cls._make_user_sessions_key(user_id)
            
            # Удаляем токен из Redis в обоих форматах
            pipeline = redis.pipeline()
            pipeline.delete(token_key)
            pipeline.delete(session_key)
            pipeline.srem(user_sessions_key, token)
            results = await pipeline.execute()
            
            success = any(results)
            if success:
                logger.info(f"[TokenStorage.invalidate_token] Токен успешно инвалидирован для пользователя {user_id}")
            else:
                logger.warning(f"[TokenStorage.invalidate_token] Токен не найден: {token_key}")
                
            return success
            
        except Exception as e:
            logger.error(f"[TokenStorage.invalidate_token] Ошибка при инвалидации токена: {e}")
            return False

    @classmethod
    async def invalidate_all_tokens(cls, user_id: str) -> int:
        """
        Инвалидирует все токены пользователя

        Args:
            user_id: ID пользователя

        Returns:
            int: Количество инвалидированных токенов
        """
        try:
            # Получаем список сессий пользователя
            user_sessions_key = cls._make_user_sessions_key(user_id)
            tokens = await redis.smembers(user_sessions_key)
            
            if not tokens:
                logger.warning(f"[TokenStorage.invalidate_all_tokens] Нет активных сессий пользователя {user_id}")
                return 0
                
            count = 0
            for token in tokens:
                # Декодируем JWT токен
                try:
                    payload = JWTCodec.decode(token)
                    if payload:
                        username = payload.username
                        
                        # Формируем ключи для Redis
                        token_key = cls._make_token_key(user_id, username, token)
                        session_key = cls._make_session_key(user_id, token)
                        
                        # Удаляем токен из Redis
                        pipeline = redis.pipeline()
                        pipeline.delete(token_key)
                        pipeline.delete(session_key)
                        results = await pipeline.execute()
                        
                        count += 1
                except Exception as e:
                    logger.error(f"[TokenStorage.invalidate_all_tokens] Ошибка при обработке токена: {e}")
                    continue
            
            # Удаляем список сессий пользователя
            await redis.delete(user_sessions_key)
            
            logger.info(f"[TokenStorage.invalidate_all_tokens] Инвалидировано {count} токенов пользователя {user_id}")
            return count
            
        except Exception as e:
            logger.error(f"[TokenStorage.invalidate_all_tokens] Ошибка при инвалидации всех токенов: {e}")
            return 0

    @classmethod
    async def get_session_data(cls, token: str) -> Optional[Dict[str, Any]]:
        """
        Получает данные сессии

        Args:
            token: JWT токен

        Returns:
            Dict[str, Any]: Данные сессии или None
        """
        valid, data = await cls.validate_token(token)
        return data if valid else None

    @staticmethod
    async def get(token_key: str) -> Optional[str]:
        """
        Получает токен из хранилища.

        Args:
            token_key: Ключ токена

        Returns:
            str или None, если токен не найден
        """
        logger.debug(f"[tokenstorage.get] Запрос токена: {token_key}")
        return await redis.get(token_key)

    @staticmethod
    async def exists(token_key: str) -> bool:
        """
        Проверяет наличие токена в хранилище.

        Args:
            token_key: Ключ токена

        Returns:
            bool: True, если токен существует
        """
        return bool(await redis.execute("EXISTS", token_key))

    @staticmethod
    async def save_token(token_key: str, data: Dict[str, Any], life_span: int) -> bool:
        """
        Сохраняет токен в хранилище с указанным временем жизни.

        Args:
            token_key: Ключ токена
            data: Данные токена
            life_span: Время жизни токена в секундах

        Returns:
            bool: True, если токен успешно сохранен
        """
        try:
            # Если данные не строка, преобразуем их в JSON
            value = json.dumps(data) if isinstance(data, dict) else data

            # Сохраняем токен и устанавливаем время жизни
            await redis.set(token_key, value, ex=life_span)

            return True
        except Exception as e:
            logger.error(f"[tokenstorage.save_token] Ошибка сохранения токена: {str(e)}")
            return False

    @staticmethod
    async def create_onetime(user: AuthInput) -> str:
        """
        Создает одноразовый токен для пользователя.

        Args:
            user: Объект пользователя

        Returns:
            str: Сгенерированный токен
        """
        life_span = ONETIME_TOKEN_LIFE_SPAN
        exp = datetime.now(tz=timezone.utc) + timedelta(seconds=life_span)
        one_time_token = JWTCodec.encode(user, exp)

        # Сохраняем токен в Redis
        token_key = f"{user.id}-{user.username}-{one_time_token}"
        await TokenStorage.save_token(token_key, "TRUE", life_span)

        return one_time_token

    @staticmethod
    async def revoke(token: str) -> bool:
        """
        Отзывает токен.

        Args:
            token: Токен для отзыва

        Returns:
            bool: True, если токен успешно отозван
        """
        try:
            logger.debug("[tokenstorage.revoke] Отзыв токена")

            # Декодируем токен
            payload = JWTCodec.decode(token)
            if not payload:
                logger.warning("[tokenstorage.revoke] Невозможно декодировать токен")
                return False

            # Формируем ключи
            token_key = f"{payload.user_id}-{payload.username}-{token}"
            user_sessions_key = f"user_sessions:{payload.user_id}"

            # Удаляем токен и запись из списка сессий пользователя
            pipe = redis.pipeline()
            await pipe.delete(token_key)
            await pipe.srem(user_sessions_key, token)
            await pipe.execute()

            return True
        except Exception as e:
            logger.error(f"[tokenstorage.revoke] Ошибка отзыва токена: {str(e)}")
            return False

    @staticmethod
    async def revoke_all(user: AuthInput) -> bool:
        """
        Отзывает все токены пользователя.

        Args:
            user: Объект пользователя

        Returns:
            bool: True, если все токены успешно отозваны
        """
        try:
            # Формируем ключи
            user_sessions_key = f"user_sessions:{user.id}"

            # Получаем все токены пользователя
            tokens = await redis.smembers(user_sessions_key)
            if not tokens:
                return True

            # Формируем список ключей для удаления
            keys_to_delete = [f"{user.id}-{user.username}-{token}" for token in tokens]
            keys_to_delete.append(user_sessions_key)

            # Удаляем все токены и список сессий
            await redis.delete(*keys_to_delete)

            return True
        except Exception as e:
            logger.error(f"[tokenstorage.revoke_all] Ошибка отзыва всех токенов: {str(e)}")
            return False