gemini_engine.py

"""
Gemini AI Engine Module
Handles semantic email analysis using Google Gemini API
Provides mock mode for development without live API calls
"""

import os
import re
import json
from datetime import datetime, timezone
from typing import Dict, List, Optional
from dotenv import load_dotenv
import google.generativeai as genai

# Load environment variables from .env file (override=True to reload on changes)
load_dotenv(override=True)

# Configuration - Global variables that can be modified
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', '').strip()
MOCK_MODE = os.getenv('GEMINI_MOCK_MODE', 'true').lower() == 'true'

# Initialize Gemini only if API key is present and not in mock mode
if GEMINI_API_KEY and not MOCK_MODE:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        # Use gemini-2.0-flash (latest stable flash model)
        model = genai.GenerativeModel('gemini-2.0-flash')
        print("✅ Gemini API initialized successfully with gemini-2.0-flash")
    except Exception as e:
        print(f"Warning: Failed to initialize Gemini API: {e}")
        print("Please verify:")
        print("1. API key is valid at https://aistudio.google.com/app/apikey")
        print("2. Generative Language API is enabled in your Google Cloud project")
        model = None
        MOCK_MODE = True
else:
    model = None


def set_mock_mode(enabled: bool):
    """
    Dynamically set mock mode at runtime
    
    Args:
        enabled: True for mock mode, False for live mode
    """
    global MOCK_MODE, model
    MOCK_MODE = enabled
    
    # Reinitialize model if switching to live mode
    if not MOCK_MODE and GEMINI_API_KEY:
        try:
            genai.configure(api_key=GEMINI_API_KEY)
            model = genai.GenerativeModel('gemini-2.0-flash')
            print("✅ Switched to Live mode with gemini-2.0-flash")
        except Exception as e:
            print(f"Error initializing Gemini: {e}")
            model = None
            MOCK_MODE = True
    else:
        model = None
        print("ℹ️ Switched to Mock mode")


def redact_urls(text: str) -> tuple[str, List[str]]:
    """
    Redact live URLs from text, extracting only registrable domains
    
    Args:
        text: Input text containing URLs
        
    Returns:
        Tuple of (redacted_text, list_of_domains)
    """
    if not text:
        return "", []
    
    domains = []
    
    # Pattern to match URLs
    url_pattern = r'https?://[^\s<>"{}|\\^`\[\]]+'
    
    def extract_domain(match):
        url = match.group(0)
        # Extract domain from URL
        domain_match = re.search(r'https?://([^/:?#]+)', url)
        if domain_match:
            domain = domain_match.group(1)
            # Remove 'www.' prefix if present
            domain = re.sub(r'^www\.', '', domain)
            domains.append(domain)
            return f"[DOMAIN:{domain}]"
        return "[REDACTED_URL]"
    
    redacted = re.sub(url_pattern, extract_domain, text)
    return redacted, domains


def redact_sensitive_info(text: str) -> str:
    """
    Redact sensitive information like emails, phone numbers, etc.
    
    Args:
        text: Input text
        
    Returns:
        Redacted text
    """
    if not text:
        return ""
    
    # Redact email addresses
    text = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL]', text)
    
    # Redact phone numbers (various formats)
    text = re.sub(r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE]', text)
    text = re.sub(r'\+\d{1,3}[-.\s]?\d{1,4}[-.\s]?\d{1,4}[-.\s]?\d{1,9}', '[PHONE]', text)
    
    # Redact credit card patterns
    text = re.sub(r'\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b', '[CARD]', text)
    
    # Redact SSN patterns
    text = re.sub(r'\b\d{3}-\d{2}-\d{4}\b', '[SSN]', text)
    
    return text


def prepare_email_for_analysis(parsed_email: Dict, heuristics: List[str]) -> str:
    """
    Prepare email content for Gemini analysis with redaction
    
    Args:
        parsed_email: Dictionary containing email data
        heuristics: List of detected heuristic indicators
        
    Returns:
        Formatted and redacted text for analysis
    """
    subject = parsed_email.get('subject', '(No Subject)')
    sender = parsed_email.get('from', parsed_email.get('sender', 'Unknown'))
    body = parsed_email.get('body_text', '')
    attachments = parsed_email.get('attachments', [])
    
    # Redact URLs and extract domains
    redacted_body, domains = redact_urls(body)
    
    # Redact other sensitive info
    redacted_body = redact_sensitive_info(redacted_body)
    
    # Build analysis prompt
    prompt = f"""Analyze this email for phishing threats and provide a detailed risk assessment.

EMAIL DATA:
- Subject: {subject}
- From: {sender}
- Has Attachments: {len(attachments) > 0}
- Attachment Types: {', '.join(attachments) if attachments else 'None'}
- Domains Found: {', '.join(domains) if domains else 'None'}

DETECTED HEURISTIC INDICATORS:
{', '.join(heuristics) if heuristics else 'None'}

EMAIL CONTENT (redacted):
{redacted_body[:3000]}

ANALYSIS REQUIRED:
1. Risk Score (0-100): Assess the overall phishing risk
2. Intents: What is the sender trying to achieve? (e.g., credential_theft, financial_fraud, malware_delivery)
3. Indicators: Specific phishing indicators found (e.g., urgency_tactics, impersonation, suspicious_links)
4. Safe Summary: A brief, safe-to-display summary of the email
5. Recommendations: Specific actions the user should take

Provide your response in valid JSON format with these exact keys:
{{
  "risk_score": <number 0-100>,
  "intents": [<list of intent strings>],
  "indicators": [<list of indicator strings>],
  "safe_summary": "<summary string>",
  "recommendations": [<list of recommendation strings>]
}}"""
    
    return prompt


def parse_gemini_response(response_text: str) -> Dict:
    """
    Parse Gemini response and extract structured data
    
    Args:
        response_text: Raw response from Gemini
        
    Returns:
        Structured dictionary
    """
    try:
        # Try to extract JSON from response
        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
        if json_match:
            data = json.loads(json_match.group(0))
        else:
            # Fallback parsing if no JSON found
            data = {
                "risk_score": 50.0,
                "intents": ["unknown"],
                "indicators": ["unable_to_parse"],
                "safe_summary": "Analysis incomplete",
                "recommendations": ["Manual review recommended"]
            }
        
        # Validate and clamp risk_score
        risk_score = float(data.get('risk_score', 50.0))
        risk_score = max(0.0, min(100.0, risk_score))
        data['risk_score'] = risk_score
        
        # Ensure all required fields exist
        data.setdefault('intents', [])
        data.setdefault('indicators', [])
        data.setdefault('safe_summary', '')
        data.setdefault('recommendations', [])
        
        return data
        
    except (json.JSONDecodeError, ValueError) as e:
        # Return safe fallback
        return {
            "risk_score": 50.0,
            "intents": ["parse_error"],
            "indicators": ["response_parsing_failed"],
            "safe_summary": "Unable to parse AI response",
            "recommendations": ["Manual review required"],
            "error": str(e)
        }


def generate_mock_analysis(parsed_email: Dict, heuristics: List[str]) -> Dict:
    """
    Generate mock analysis data for development/testing
    
    Args:
        parsed_email: Email data dictionary
        heuristics: List of heuristic indicators
        
    Returns:
        Mock analysis dictionary
    """
    # Calculate mock risk score based on heuristics
    base_score = len(heuristics) * 15
    
    # Add randomness based on content
    subject = parsed_email.get('subject', '').lower()
    body = parsed_email.get('body_text', '').lower()
    
    risk_modifiers = 0
    if 'urgent' in subject or 'urgent' in body:
        risk_modifiers += 10
    if 'verify' in subject or 'account' in subject:
        risk_modifiers += 10
    if 'click' in body or 'link' in body:
        risk_modifiers += 5
    
    risk_score = min(100.0, max(0.0, float(base_score + risk_modifiers)))
    
    # Determine intents based on heuristics
    intents = []
    if 'credential' in ' '.join(heuristics).lower():
        intents.append('credential_theft')
    if 'payment' in ' '.join(heuristics).lower() or 'urgency' in ' '.join(heuristics).lower():
        intents.append('financial_fraud')
    if 'attachment' in ' '.join(heuristics).lower():
        intents.append('malware_delivery')
    if not intents:
        intents = ['information_gathering']
    
    # Build indicators (deduplicate as we go)
    indicators = []
    for h in heuristics:
        if 'urgency' in h.lower() and 'urgency_tactics' not in indicators:
            indicators.append('urgency_tactics')
        if 'credential' in h.lower() and 'credential_request' not in indicators:
            indicators.append('credential_request')
        if 'domain' in h.lower() and 'suspicious_domain' not in indicators:
            indicators.append('suspicious_domain')
        if 'link' in h.lower() and 'suspicious_links' not in indicators:
            indicators.append('suspicious_links')
    
    if not indicators:
        indicators = ['low_confidence_signals']
    
    # Generate recommendations (avoid duplicates)
    recommendations = []
    if risk_score >= 70:
        recommendations = [
            "Do not click any links in this email",
            "Do not provide any personal information",
            "Mark as spam and delete",
            "Report to your IT security team"
        ]
    elif risk_score >= 30:
        recommendations = [
            "Verify sender identity through alternative channel",
            "Examine links carefully before clicking",
            "Be cautious with any requests"
        ]
    else:
        recommendations = [
            "Email appears safe but remain vigilant",
            "Verify any unexpected requests independently"
        ]
    
    # Generate safe summary
    subject_clean = parsed_email.get('subject', 'No subject')[:100]
    safe_summary = f"Email regarding: {subject_clean}. Risk level: {'HIGH' if risk_score >= 70 else 'MEDIUM' if risk_score >= 30 else 'LOW'}."
    
    return {
        "risk_score": risk_score,
        "intents": intents,
        "indicators": indicators,
        "safe_summary": safe_summary,
        "recommendations": recommendations,
        "processed_at": datetime.now(timezone.utc).isoformat(),
        "mock_mode": True
    }


def analyze_email(parsed_email: Dict, heuristics: List[str], mode: Optional[str] = None) -> Dict:
    """
    Analyze email using Gemini AI for semantic phishing detection
    
    Args:
        parsed_email: Dictionary containing email data with keys:
            - subject: Email subject
            - from/sender: Sender address
            - date: Date sent
            - body_text: Email body content
            - attachments: List of attachment names
        heuristics: List of detected heuristic indicators
        mode: Override for mock/live mode ('mock' or 'live')
        
    Returns:
        Dictionary containing:
            - risk_score: float (0-100)
            - intents: list[str] - What the sender is trying to achieve
            - indicators: list[str] - Specific phishing indicators found
            - safe_summary: str - Safe-to-display summary
            - recommendations: list[str] - Recommended actions
            - processed_at: ISO8601 UTC timestamp
    """
    # Determine mode
    use_mock = MOCK_MODE
    if mode == 'mock':
        use_mock = True
    elif mode == 'live':
        use_mock = False
    
    # Use mock mode if no API key or explicitly requested
    if use_mock or not GEMINI_API_KEY or model is None:
        return generate_mock_analysis(parsed_email, heuristics)
    
    try:
        # Prepare email for analysis
        prompt = prepare_email_for_analysis(parsed_email, heuristics)
        
        # Call Gemini API
        response = model.generate_content(prompt)
        
        # Parse response
        analysis = parse_gemini_response(response.text)
        
        # Add metadata
        analysis['processed_at'] = datetime.now(timezone.utc).isoformat()
        analysis['mock_mode'] = False
        
        return analysis
        
    except Exception as e:
        # Fallback to mock on error
        print(f"Gemini API error: {e}. Falling back to mock mode.")
        result = generate_mock_analysis(parsed_email, heuristics)
        result['error'] = str(e)
        result['fallback'] = True
        return result


def batch_analyze_emails(emails: List[tuple], mode: Optional[str] = None) -> List[Dict]:
    """
    Analyze multiple emails in batch
    
    Args:
        emails: List of tuples (parsed_email, heuristics)
        mode: Override for mock/live mode
        
    Returns:
        List of analysis dictionaries
    """
    results = []
    for parsed_email, heuristics in emails:
        analysis = analyze_email(parsed_email, heuristics, mode)
        results.append(analysis)
    return results


# Configuration helper
def set_mode(mock: bool):
    """
    Set the analysis mode globally
    
    Args:
        mock: True for mock mode, False for live API mode
    """
    global MOCK_MODE
    MOCK_MODE = mock


def is_mock_mode() -> bool:
    """
    Check if currently in mock mode
    
    Returns:
        True if mock mode, False if live mode
    """
    return MOCK_MODE or not GEMINI_API_KEY or model is None