Gondulf/src/gondulf/storage.py

"""
In-memory storage for short-lived codes with TTL.

Provides simple dict-based storage for email verification codes and authorization
codes with automatic expiration checking on access.
"""

import json
import logging
import time
from typing import Union

logger = logging.getLogger("gondulf.storage")


class CodeStore:
    """
    In-memory storage for domain verification codes with TTL.

    Stores codes with expiration timestamps and automatically removes expired
    codes on access. No background cleanup needed - cleanup happens lazily.
    """

    def __init__(self, ttl_seconds: int = 600):
        """
        Initialize code store.

        Args:
            ttl_seconds: Time-to-live for codes in seconds (default: 600 = 10 minutes)
        """
        self._store: dict[str, tuple[Union[str, dict], float]] = {}
        self._ttl = ttl_seconds
        logger.debug(f"CodeStore initialized with TTL={ttl_seconds}s")

    def store(self, key: str, value: Union[str, dict], ttl: int | None = None) -> None:
        """
        Store value (string or dict) with expiry timestamp.

        Args:
            key: Storage key (typically email address or code identifier)
            value: Value to store (string for simple codes, dict for authorization code metadata)
            ttl: Optional TTL override in seconds (default: use instance TTL)
        """
        expiry = time.time() + (ttl if ttl is not None else self._ttl)
        self._store[key] = (value, expiry)
        logger.debug(f"Value stored for key={key} expires_in={ttl if ttl is not None else self._ttl}s")

    def verify(self, key: str, code: str) -> bool:
        """
        Verify code matches stored value and remove from store.

        Checks both expiration and code matching. If valid, removes the code
        from storage (single-use). Expired codes are also removed.

        Args:
            key: Storage key to verify
            code: Code to verify

        Returns:
            True if code matches and is not expired, False otherwise
        """
        if key not in self._store:
            logger.debug(f"Verification failed: key={key} not found")
            return False

        stored_code, expiry = self._store[key]

        # Check expiration
        if time.time() > expiry:
            del self._store[key]
            logger.debug(f"Verification failed: key={key} expired")
            return False

        # Check code match
        if code != stored_code:
            logger.debug(f"Verification failed: key={key} code mismatch")
            return False

        # Valid - remove from store (single use)
        del self._store[key]
        logger.info(f"Code verified successfully for key={key}")
        return True

    def get(self, key: str) -> Union[str, dict, None]:
        """
        Get value without removing it.

        Checks expiration and removes expired values.

        Args:
            key: Storage key to retrieve

        Returns:
            Value (str or dict) if exists and not expired, None otherwise
        """
        if key not in self._store:
            return None

        stored_value, expiry = self._store[key]

        # Check expiration
        if time.time() > expiry:
            del self._store[key]
            return None

        return stored_value

    def delete(self, key: str) -> None:
        """
        Explicitly delete a code from storage.

        Args:
            key: Storage key to delete
        """
        if key in self._store:
            del self._store[key]
            logger.debug(f"Code deleted for key={key}")

    def cleanup_expired(self) -> int:
        """
        Manually cleanup all expired codes.

        This is optional - cleanup happens automatically on access. But can be
        called periodically if needed to free memory.

        Returns:
            Number of expired codes removed
        """
        now = time.time()
        expired_keys = [key for key, (_, expiry) in self._store.items() if now > expiry]

        for key in expired_keys:
            del self._store[key]

        if expired_keys:
            logger.debug(f"Cleaned up {len(expired_keys)} expired codes")

        return len(expired_keys)

    def size(self) -> int:
        """
        Get number of codes currently in storage (including expired).

        Returns:
            Number of codes in storage
        """
        return len(self._store)

    def clear(self) -> None:
        """Clear all codes from storage."""
        self._store.clear()
        logger.debug("Code store cleared")