mongo_client.py

"""
PhishGuard M4 - MongoDB Client Module
Secure connection handling for MongoDB Atlas
Production-grade database connectivity
"""

import os
from typing import Optional
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global MongoDB client instance
_mongo_client: Optional[MongoClient] = None


def get_mongo_client(mongo_uri: str, force_new: bool = False) -> MongoClient:
    """
    Get or create a MongoDB client instance.
    
    Uses connection pooling and maintains a singleton client instance
    for the application lifecycle.
    
    Args:
        mongo_uri: MongoDB Atlas connection URI from decrypted secrets
        force_new: If True, creates a new client instance (for testing)
        
    Returns:
        MongoClient instance ready for database operations
        
    Raises:
        ConnectionFailure: If connection to MongoDB fails
        ValueError: If mongo_uri is invalid or empty
        
    Example:
        >>> from crypto_simple import load_encrypted
        >>> secrets = load_encrypted('secrets.enc', passphrase)
        >>> client = get_mongo_client(secrets['MONGO_URI'])
        >>> db = client['phishguard']
    """
    global _mongo_client
    
    # Validate URI
    if not mongo_uri or not isinstance(mongo_uri, str):
        raise ValueError("MongoDB URI must be a non-empty string")
    
    if not mongo_uri.startswith(('mongodb://', 'mongodb+srv://')):
        raise ValueError("Invalid MongoDB URI format. Must start with 'mongodb://' or 'mongodb+srv://'")
    
    # Return existing client if available
    if _mongo_client is not None and not force_new:
        try:
            # Test connection
            _mongo_client.admin.command('ping')
            logger.info("✅ Reusing existing MongoDB connection")
            return _mongo_client
        except Exception as e:
            logger.warning(f"⚠️ Existing MongoDB connection invalid, creating new: {e}")
            _mongo_client = None
    
    # Create new client
    try:
        logger.info("🔌 Connecting to MongoDB Atlas...")
        
        # Connection settings optimized for security and reliability
        client = MongoClient(
            mongo_uri,
            serverSelectionTimeoutMS=5000,  # 5 second timeout
            connectTimeoutMS=10000,         # 10 second connection timeout
            socketTimeoutMS=30000,          # 30 second socket timeout
            maxPoolSize=10,                 # Connection pool size
            minPoolSize=1,                  # Minimum connections
            retryWrites=True,               # Retry write operations
            retryReads=True,                # Retry read operations
            tls=True,                       # Enforce TLS/SSL
            tlsAllowInvalidCertificates=False,  # Strict certificate validation
        )
        
        # Verify connection
        client.admin.command('ping')
        
        logger.info("✅ MongoDB connection established successfully")
        
        # Cache the client
        _mongo_client = client
        
        return client
        
    except ConnectionFailure as e:
        logger.error(f"❌ MongoDB connection failed: {e}")
        raise ConnectionFailure(
            f"Failed to connect to MongoDB. Check your MONGO_URI and network connectivity: {e}"
        )
    except ServerSelectionTimeoutError as e:
        logger.error(f"❌ MongoDB server selection timeout: {e}")
        raise ConnectionFailure(
            f"MongoDB server selection timeout. Verify URI and network access: {e}"
        )
    except Exception as e:
        logger.error(f"❌ Unexpected error connecting to MongoDB: {e}")
        raise


def close_mongo_client() -> None:
    """
    Close the global MongoDB client connection.
    
    Should be called on application shutdown for graceful cleanup.
    """
    global _mongo_client
    
    if _mongo_client is not None:
        try:
            _mongo_client.close()
            logger.info("🔌 MongoDB connection closed")
        except Exception as e:
            logger.warning(f"⚠️ Error closing MongoDB connection: {e}")
        finally:
            _mongo_client = None


def test_connection(mongo_uri: str) -> bool:
    """
    Test MongoDB connection without caching the client.
    
    Useful for validation during setup or configuration.
    
    Args:
        mongo_uri: MongoDB Atlas connection URI
        
    Returns:
        True if connection succeeds, False otherwise
    """
    try:
        client = MongoClient(
            mongo_uri,
            serverSelectionTimeoutMS=5000
        )
        client.admin.command('ping')
        client.close()
        return True
    except Exception as e:
        logger.error(f"❌ Connection test failed: {e}")
        return False


# Security notes:
# - MONGO_URI should always be loaded from encrypted secrets (M3)
# - Never hardcode credentials or URIs in source code
# - Use MongoDB Atlas IP allowlisting for additional security
# - Rotate database passwords regularly
# - Enable MongoDB Atlas encryption at rest