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 logging
import time
from typing import Dict, Optional, Tuple

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[str, float]] = {}
        self._ttl = ttl_seconds
        logger.debug(f"CodeStore initialized with TTL={ttl_seconds}s")

    def store(self, key: str, code: str) -> None:
        """
        Store verification code with expiry timestamp.

        Args:
            key: Storage key (typically email address or similar identifier)
            code: Verification code to store
        """
        expiry = time.time() + self._ttl
        self._store[key] = (code, expiry)
        logger.debug(f"Code stored for key={key} expires_in={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) -> Optional[str]:
        """
        Get code without removing it (for testing/debugging).

        Checks expiration and removes expired codes.

        Args:
            key: Storage key to retrieve

        Returns:
            Code if exists and not expired, None otherwise
        """
        if key not in self._store:
            return None

        stored_code, expiry = self._store[key]

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

        return stored_code

    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")