storage.py

"""
PhishGuard M4 - Storage Module
MongoDB operations for email analysis persistence
Integrates with M3 signing for data integrity verification
"""

from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional
from pymongo import MongoClient, ASCENDING, DESCENDING
from pymongo.errors import DuplicateKeyError, PyMongoError
import logging

from mongo_client import get_mongo_client
from signing import verify
from crypto_simple import load_encrypted

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

# Database and collection names
DATABASE_NAME = "phishguard"
COLLECTION_NAME = "email_analyses"


def _get_collection(mongo_client: MongoClient):
    """
    Get the analyses collection with proper indexing.
    
    Args:
        mongo_client: Active MongoDB client instance
        
    Returns:
        MongoDB collection object
    """
    db = mongo_client[DATABASE_NAME]
    collection = db[COLLECTION_NAME]
    
    # Create indexes if they don't exist
    try:
        # Unique index on gmail_id to prevent duplicates
        collection.create_index(
            [("gmail_id", ASCENDING)],
            unique=True,
            name="gmail_id_unique"
        )
        logger.info("✅ Ensured unique index on gmail_id")
        
        # Index on processed_at for time-based queries
        collection.create_index(
            [("processed_at", DESCENDING)],
            name="processed_at_desc"
        )
        logger.info("✅ Ensured index on processed_at")
        
        # Index on risk_label for filtering
        collection.create_index(
            [("risk_label", ASCENDING)],
            name="risk_label_asc"
        )
        logger.info("✅ Ensured index on risk_label")
        
        # Optional: TTL index to auto-delete old records after 7 days
        # Uncomment if you want automatic cleanup
        # collection.create_index(
        #     [("processed_at", ASCENDING)],
        #     expireAfterSeconds=7 * 24 * 60 * 60,  # 7 days
        #     name="processed_at_ttl"
        # )
        # logger.info("✅ Ensured TTL index (7 days retention)")
        
    except Exception as e:
        logger.warning(f"⚠️ Index creation warning: {e}")
    
    return collection


def insert_analysis(
    analysis_doc: Dict[str, Any],
    app_start_time: datetime,
    mongo_uri: str,
    signing_secret: bytes
) -> bool:
    """
    Insert a verified analysis document into MongoDB.
    
    Only inserts analyses processed during the current app session
    (after app_start_time) and with valid signatures.
    
    Args:
        analysis_doc: Analysis document with signature
        app_start_time: App session start time (for session filtering)
        mongo_uri: MongoDB connection URI from secrets
        signing_secret: Secret for signature verification
        
    Returns:
        True if inserted successfully, False otherwise
        
    Example:
        >>> from datetime import datetime
        >>> app_start = datetime.now()
        >>> analysis = {
        ...     'gmail_id': '18ab123...',
        ...     'sender': 'test@example.com',
        ...     'subject': 'Test Email',
        ...     'risk_score': 85,
        ...     'risk_label': 'HIGH_RISK',
        ...     'processed_at': datetime.now(),
        ...     'signature': 'base64sig...'
        ... }
        >>> insert_analysis(analysis, app_start, mongo_uri, secret)
    """
    try:
        # Validate required fields
        required_fields = ['gmail_id', 'processed_at', 'signature']
        for field in required_fields:
            if field not in analysis_doc:
                logger.error(f"❌ Missing required field: {field}")
                return False
        
        # Check if analysis is from current session
        processed_at = analysis_doc.get('processed_at')
        if isinstance(processed_at, str):
            processed_at = datetime.fromisoformat(processed_at)
        
        if processed_at < app_start_time:
            logger.info(f"⏭️ Skipping analysis from before app start: {analysis_doc.get('gmail_id')}")
            return False
        
        # Verify signature before insertion
        # Note: For verification, we need to ensure the document structure matches what was signed
        # If processed_at is a datetime object, convert to string for verification
        signature = analysis_doc.get('signature')
        verification_doc = analysis_doc.copy()
        if isinstance(verification_doc.get('processed_at'), datetime):
            verification_doc['processed_at'] = verification_doc['processed_at'].isoformat()
        
        is_valid = verify(verification_doc, signature, signing_secret)
        
        if not is_valid:
            logger.error(f"❌ Signature verification failed for: {analysis_doc.get('gmail_id')}")
            return False
        
        # Get MongoDB client and collection
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        # Insert document
        result = collection.insert_one(analysis_doc)
        
        if result.inserted_id:
            logger.info(f"✅ Inserted analysis: {analysis_doc.get('gmail_id')} (Risk: {analysis_doc.get('risk_label')})")
            return True
        else:
            logger.warning(f"⚠️ Insert returned no ID for: {analysis_doc.get('gmail_id')}")
            return False
            
    except DuplicateKeyError:
        logger.info(f"⏭️ Duplicate gmail_id, skipping: {analysis_doc.get('gmail_id')}")
        return False
    except PyMongoError as e:
        logger.error(f"❌ MongoDB error during insert: {e}")
        return False
    except Exception as e:
        logger.error(f"❌ Unexpected error during insert: {e}")
        return False


def load_analyses(
    mongo_uri: str,
    signing_secret: bytes,
    filter_by: Optional[Dict[str, Any]] = None,
    limit: Optional[int] = None,
    skip: int = 0
) -> List[Dict[str, Any]]:
    """
    Load and verify analysis documents from MongoDB.
    
    Retrieves documents, verifies signatures, and returns only valid records.
    Invalid signatures are logged and excluded from results.
    
    Args:
        mongo_uri: MongoDB connection URI from secrets
        signing_secret: Secret for signature verification
        filter_by: MongoDB query filter (e.g., {'risk_label': 'HIGH_RISK'})
        limit: Maximum number of documents to return
        skip: Number of documents to skip (for pagination)
        
    Returns:
        List of verified analysis documents with added 'signature_valid' field
        
    Example:
        >>> # Load all analyses
        >>> analyses = load_analyses(mongo_uri, secret)
        >>> 
        >>> # Load only high-risk analyses
        >>> high_risk = load_analyses(
        ...     mongo_uri, secret,
        ...     filter_by={'risk_label': 'HIGH_RISK'}
        ... )
        >>> 
        >>> # Paginated loading
        >>> page_1 = load_analyses(mongo_uri, secret, limit=10, skip=0)
        >>> page_2 = load_analyses(mongo_uri, secret, limit=10, skip=10)
    """
    try:
        # Get MongoDB client and collection
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        # Build query
        query = filter_by if filter_by else {}
        
        # Fetch documents
        cursor = collection.find(query).sort("processed_at", DESCENDING)
        
        if skip > 0:
            cursor = cursor.skip(skip)
        if limit:
            cursor = cursor.limit(limit)
        
        documents = list(cursor)
        
        logger.info(f"📥 Fetched {len(documents)} document(s) from MongoDB")
        
        # Verify signatures
        verified_docs = []
        invalid_count = 0
        
        for doc in documents:
            signature = doc.get('signature')
            
            if not signature:
                logger.warning(f"⚠️ Document missing signature: {doc.get('gmail_id')}")
                doc['signature_valid'] = False
                invalid_count += 1
                verified_docs.append(doc)
                continue
            
            # Check if manually verified (takes precedence)
            if doc.get('signature_manually_verified'):
                doc['signature_valid'] = True
                verified_docs.append(doc)
                continue
            
            # Verify signature - convert datetime to ISO string for verification
            verification_doc = doc.copy()
            if isinstance(verification_doc.get('processed_at'), datetime):
                verification_doc['processed_at'] = verification_doc['processed_at'].isoformat()
            
            is_valid = verify(verification_doc, signature, signing_secret)
            doc['signature_valid'] = is_valid
            
            if not is_valid:
                logger.warning(f"⚠️ Invalid signature for: {doc.get('gmail_id')}")
                invalid_count += 1
            
            verified_docs.append(doc)
        
        if invalid_count > 0:
            logger.warning(f"⚠️ Found {invalid_count} document(s) with invalid signatures")
        else:
            logger.info(f"✅ All {len(verified_docs)} document(s) have valid signatures")
        
        return verified_docs
        
    except PyMongoError as e:
        logger.error(f"❌ MongoDB error during load: {e}")
        return []
    except Exception as e:
        logger.error(f"❌ Unexpected error during load: {e}")
        return []


def get_analysis_by_gmail_id(
    gmail_id: str,
    mongo_uri: str,
    signing_secret: bytes
) -> Optional[Dict[str, Any]]:
    """
    Retrieve a specific analysis by Gmail ID.
    
    Args:
        gmail_id: Gmail message ID
        mongo_uri: MongoDB connection URI
        signing_secret: Secret for signature verification
        
    Returns:
        Verified analysis document or None if not found
    """
    try:
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        doc = collection.find_one({"gmail_id": gmail_id})
        
        if not doc:
            logger.info(f"📭 No analysis found for gmail_id: {gmail_id}")
            return None
        
        # Verify signature
        signature = doc.get('signature')
        if signature:
            doc['signature_valid'] = verify(doc, signature, signing_secret)
        else:
            doc['signature_valid'] = False
        
        return doc
        
    except Exception as e:
        logger.error(f"❌ Error retrieving analysis: {e}")
        return None


def count_analyses(
    mongo_uri: str,
    filter_by: Optional[Dict[str, Any]] = None
) -> int:
    """
    Count analysis documents matching filter.
    
    Args:
        mongo_uri: MongoDB connection URI
        filter_by: MongoDB query filter
        
    Returns:
        Count of matching documents
    """
    try:
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        query = filter_by if filter_by else {}
        count = collection.count_documents(query)
        
        return count
        
    except Exception as e:
        logger.error(f"❌ Error counting analyses: {e}")
        return 0


def delete_old_analyses(
    mongo_uri: str,
    days_old: int = 7
) -> int:
    """
    Delete analyses older than specified days.
    
    Args:
        mongo_uri: MongoDB connection URI
        days_old: Delete analyses older than this many days
        
    Returns:
        Number of documents deleted
    """
    try:
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        cutoff_date = datetime.now() - timedelta(days=days_old)
        
        result = collection.delete_many({"processed_at": {"$lt": cutoff_date}})
        
        deleted_count = result.deleted_count
        logger.info(f"🗑️ Deleted {deleted_count} analysis document(s) older than {days_old} days")
        
        return deleted_count
        
    except Exception as e:
        logger.error(f"❌ Error deleting old analyses: {e}")
        return 0


def mark_signature_as_valid(
    gmail_id: str,
    mongo_uri: str
) -> bool:
    """
    Mark an analysis document's signature as manually verified.
    
    This is used when a user manually verifies a signature through the UI
    and we want to update the database to reflect the verified status.
    
    Args:
        gmail_id: Gmail message ID of the analysis to update
        mongo_uri: MongoDB connection URI
        
    Returns:
        True if updated successfully, False otherwise
    """
    try:
        client = get_mongo_client(mongo_uri)
        collection = _get_collection(client)
        
        # Update the document to add a manual verification flag
        result = collection.update_one(
            {"gmail_id": gmail_id},
            {
                "$set": {
                    "signature_manually_verified": True,
                    "signature_verified_at": datetime.now()
                }
            }
        )
        
        if result.matched_count > 0:
            logger.info(f"✅ Marked signature as verified for gmail_id: {gmail_id}")
            return True
        else:
            logger.warning(f"⚠️ No document found with gmail_id: {gmail_id}")
            return False
        
    except Exception as e:
        logger.error(f"❌ Error marking signature as valid: {e}")
        return False


# Security notes:
# - All documents must have valid signatures from M3
# - MONGO_URI loaded from encrypted secrets only
# - Unique index prevents duplicate gmail_id entries
# - Consider enabling TTL index for automatic data retention
# - Use filter_by carefully to prevent MongoDB injection