src/llm_agent/prompt_templates.py

"""
Prompt templates for LLM-based OSINT analysis and reporting.
This module contains various prompt templates for different types of analysis
and reporting scenarios in the OSINT early warning system.
"""

import json
from typing import Dict, List, Any, Optional
from datetime import datetime


def get_summary_prompt(entity: str, correlated_data: List[Dict], risk_assessment: Dict[str, Any]) -> str:
    """
    Generate a comprehensive analysis prompt for the main OSINT summary.
    
    Args:
        entity: The target entity being analyzed
        correlated_data: Normalized and correlated data from all sources
        risk_assessment: Risk assessment results
    
    Returns:
        Formatted prompt string for LLM analysis
    """
    
    # Format the data for better readability
    formatted_data = _format_data_for_prompt(correlated_data)
    
    prompt = f"""You are a cybersecurity intelligence analyst tasked with creating a comprehensive OSINT report. 

**ANALYSIS TARGET:** {entity}

**RISK ASSESSMENT:**
- Risk Level: {risk_assessment.get('level', 'Unknown')}
- Confidence Score: {risk_assessment.get('confidence', 0.0):.2f}
- Threat Score: {risk_assessment.get('threat_score', 'N/A')}

**INTELLIGENCE DATA:**
{formatted_data}

**INSTRUCTIONS:**
Please provide a detailed intelligence report in the following structure:

1. **EXECUTIVE SUMMARY** (2-3 sentences)
   - Brief overview of the entity and primary findings
   - Overall risk assessment conclusion

2. **KEY FINDINGS**
   - List the most significant discoveries
   - Highlight any immediate threats or concerns
   - Note correlations between different data sources

3. **THREAT INDICATORS**
   - Known malicious activities or associations
   - Vulnerability exposures
   - Suspicious patterns or behaviors

4. **TECHNICAL DETAILS**
   - Relevant technical information (IPs, domains, hashes, etc.)
   - Infrastructure details if applicable
   - Historical context from threat intelligence

5. **RISK IMPLICATIONS**
   - Potential impact to security posture
   - Attack vectors or exploitation methods
   - Business/operational risks

6. **RECOMMENDATIONS**
   - Immediate actions to take
   - Monitoring recommendations
   - Additional investigation suggestions

**FORMATTING REQUIREMENTS:**
- Use clear, professional language suitable for security teams
- Include specific indicators of compromise (IoCs) when available
- Prioritize actionable intelligence
- If data is limited, acknowledge limitations and suggest additional sources
- Use bullet points for lists and clear section headers

**CONTEXT NOTES:**
- This is an automated analysis; human verification is recommended for critical decisions
- Focus on practical, actionable intelligence over theoretical analysis
- If conflicting information exists, note the discrepancies and reliability of sources

Generate your analysis now:"""

    return prompt


def get_threat_analysis_prompt(entity: str, threat_data: List[Dict]) -> str:
    """
    Generate a focused threat analysis prompt for malware/threat-specific entities.
    
    Args:
        entity: The target entity (usually a hash, domain, or IP)
        threat_data: Threat-specific data from sources like VirusTotal
    
    Returns:
        Formatted prompt for threat-focused analysis
    """
    
    formatted_threats = _format_threat_data(threat_data)
    
    prompt = f"""You are a malware analyst examining threat intelligence data for: {entity}

**THREAT INTELLIGENCE DATA:**
{formatted_threats}

**ANALYSIS REQUIREMENTS:**
Provide a focused threat analysis covering:

1. **MALWARE CLASSIFICATION**
   - Malware family/type identification
   - Behavioral characteristics
   - Attack methodology

2. **DETECTION COVERAGE**
   - Antivirus detection rates
   - Security vendor assessments
   - False positive likelihood

3. **CAMPAIGN ANALYSIS**
   - Associated threat actors or campaigns
   - Geographic distribution
   - Timeline of activity

4. **TECHNICAL INDICATORS**
   - File properties and metadata
   - Network communications
   - System modifications

5. **MITIGATION STRATEGIES**
   - Detection rules and signatures
   - Blocking recommendations
   - Remediation steps

Focus on actionable threat intelligence for security operations teams."""

    return prompt


def get_vulnerability_analysis_prompt(entity: str, cve_data: List[Dict], asset_data: List[Dict]) -> str:
    """
    Generate a vulnerability-focused analysis prompt.
    
    Args:
        entity: The target entity (CVE, software, or system)
        cve_data: CVE and vulnerability data
        asset_data: Related asset exposure data
    
    Returns:
        Formatted prompt for vulnerability analysis
    """
    
    formatted_cves = _format_cve_data(cve_data)
    formatted_assets = _format_data_for_prompt(asset_data)
    
    prompt = f"""You are a vulnerability analyst assessing security risks for: {entity}

**VULNERABILITY DATA:**
{formatted_cves}

**ASSET EXPOSURE DATA:**
{formatted_assets}

**ANALYSIS FRAMEWORK:**
Provide a comprehensive vulnerability assessment:

1. **VULNERABILITY OVERVIEW**
   - CVE details and severity scores
   - Affected systems and versions
   - Exploitation complexity

2. **THREAT LANDSCAPE**
   - Known exploits in the wild
   - Threat actor interest
   - Attack frequency and methods

3. **ASSET IMPACT ANALYSIS**
   - Exposed systems and services
   - Business risk assessment
   - Exploitation prerequisites

4. **REMEDIATION GUIDANCE**
   - Patching recommendations
   - Compensating controls
   - Priority assessment

5. **MONITORING RECOMMENDATIONS**
   - Detection strategies
   - Log analysis guidance
   - Indicators to watch for

Focus on practical remediation guidance and risk prioritization."""

    return prompt


def get_network_analysis_prompt(entity: str, network_data: List[Dict]) -> str:
    """
    Generate a network infrastructure analysis prompt.
    
    Args:
        entity: The target entity (IP, domain, or network range)
        network_data: Network scanning and infrastructure data
    
    Returns:
        Formatted prompt for network analysis
    """
    
    formatted_network = _format_network_data(network_data)
    
    prompt = f"""You are a network security analyst examining infrastructure for: {entity}

**NETWORK INTELLIGENCE:**
{formatted_network}

**ANALYSIS OBJECTIVES:**
Provide a network security assessment covering:

1. **INFRASTRUCTURE OVERVIEW**
   - Network topology and services
   - Geographic distribution
   - Hosting providers and ASNs

2. **SERVICE ANALYSIS**
   - Open ports and services
   - Software versions and banners
   - Security configurations

3. **SECURITY POSTURE**
   - Vulnerabilities and misconfigurations
   - Security controls presence
   - Attack surface assessment

4. **THREAT INDICATORS**
   - Suspicious services or configurations
   - Known malicious associations
   - Anomalous network behavior

5. **SECURITY RECOMMENDATIONS**
   - Hardening suggestions
   - Monitoring priorities
   - Investigation actions

Focus on network security implications and defensive recommendations."""

    return prompt


def get_correlation_analysis_prompt(entity: str, multi_source_data: Dict[str, List[Dict]]) -> str:
    """
    Generate a cross-source correlation analysis prompt.
    
    Args:
        entity: The target entity
        multi_source_data: Dictionary of data from multiple sources
    
    Returns:
        Formatted prompt for correlation analysis
    """
    
    sources_summary = []
    for source, data in multi_source_data.items():
        if data:
            sources_summary.append(f"**{source.upper()}:** {len(data)} records")
            sources_summary.append(_format_data_for_prompt(data[:3]))  # Show first 3 records
    
    sources_text = "\n".join(sources_summary)
    
    prompt = f"""You are a threat intelligence analyst performing cross-source correlation analysis for: {entity}

**MULTI-SOURCE INTELLIGENCE:**
{sources_text}

**CORRELATION OBJECTIVES:**
Identify patterns, relationships, and insights across data sources:

1. **CROSS-SOURCE VALIDATION**
   - Confirm findings across multiple sources
   - Identify conflicting information
   - Assess source reliability

2. **PATTERN RECOGNITION**
   - Temporal correlations and trends
   - Geographic clustering
   - Behavioral patterns

3. **RELATIONSHIP MAPPING**
   - Associated entities and infrastructure
   - Campaign connections
   - Attribution indicators

4. **CONFIDENCE ASSESSMENT**
   - Data quality evaluation
   - Source credibility analysis
   - Confidence scoring rationale

5. **INTELLIGENCE GAPS**
   - Missing information identification
   - Additional collection requirements
   - Recommended follow-up actions

Focus on synthesizing disparate data sources into coherent intelligence assessments."""

    return prompt


def get_incident_response_prompt(entity: str, threat_data: List[Dict], urgency: str = "medium") -> str:
    """
    Generate an incident response focused prompt.
    
    Args:
        entity: The entity involved in the incident
        threat_data: Relevant threat intelligence
        urgency: Incident urgency level (low/medium/high/critical)
    
    Returns:
        Formatted prompt for incident response analysis
    """
    
    formatted_data = _format_data_for_prompt(threat_data)
    
    prompt = f"""You are an incident response analyst dealing with a {urgency.upper()} priority security incident involving: {entity}

**INCIDENT INTELLIGENCE:**
{formatted_data}

**INCIDENT RESPONSE PRIORITIES:**
Provide immediate actionable guidance:

1. **IMMEDIATE ACTIONS** (Next 1-4 hours)
   - Containment recommendations
   - Evidence preservation steps
   - Critical system protections

2. **THREAT ASSESSMENT**
   - Attacker capabilities and intent
   - Likely attack progression
   - Additional targets at risk

3. **INVESTIGATION GUIDANCE**
   - Key artifacts to collect
   - Log analysis priorities
   - Forensic examination focus

4. **COMMUNICATION REQUIREMENTS**
   - Stakeholder notification needs
   - External reporting obligations
   - Public relations considerations

5. **RECOVERY PLANNING**
   - System restoration priorities
   - Business continuity impact
   - Lessons learned capture

**URGENCY LEVEL: {urgency.upper()}**
Prioritize speed and accuracy in your recommendations. Focus on stopping ongoing threats and preserving evidence."""

    return prompt


# Helper functions for data formatting

def _format_data_for_prompt(data: List[Dict]) -> str:
    """Format data list for inclusion in prompts."""
    if not data:
        return "No data available"
    
    formatted_items = []
    for i, item in enumerate(data[:10]):  # Limit to first 10 items
        formatted_items.append(f"- Record {i+1}: {_format_dict_for_prompt(item)}")
    
    if len(data) > 10:
        formatted_items.append(f"... and {len(data) - 10} more records")
    
    return "\n".join(formatted_items)


def _format_dict_for_prompt(data: Dict) -> str:
    """Format a dictionary for readable inclusion in prompts."""
    if not data:
        return "No details available"
    
    # Prioritize important fields
    important_fields = ['type', 'source', 'confidence', 'threat_type', 'malware_family', 
                       'ip', 'domain', 'hash', 'cve_id', 'severity', 'description']
    
    formatted_parts = []
    
    # Add important fields first
    for field in important_fields:
        if field in data and data[field]:
            formatted_parts.append(f"{field}: {data[field]}")
    
    # Add other fields (limited)
    other_fields = [k for k in data.keys() if k not in important_fields][:5]
    for field in other_fields:
        if data[field]:
            formatted_parts.append(f"{field}: {str(data[field])[:100]}")
    
    return " | ".join(formatted_parts)


def _format_threat_data(threat_data: List[Dict]) -> str:
    """Format threat-specific data for prompts."""
    if not threat_data:
        return "No threat data available"
    
    formatted_items = []
    for item in threat_data:
        threat_info = []
        if 'malware_family' in item:
            threat_info.append(f"Family: {item['malware_family']}")
        if 'detection_ratio' in item:
            threat_info.append(f"Detection: {item['detection_ratio']}")
        if 'threat_type' in item:
            threat_info.append(f"Type: {item['threat_type']}")
        if 'first_seen' in item:
            threat_info.append(f"First Seen: {item['first_seen']}")
        
        formatted_items.append(" | ".join(threat_info))
    
    return "\n".join([f"- {item}" for item in formatted_items])


def _format_cve_data(cve_data: List[Dict]) -> str:
    """Format CVE-specific data for prompts."""
    if not cve_data:
        return "No CVE data available"
    
    formatted_items = []
    for cve in cve_data:
        cve_info = []
        if 'cve_id' in cve:
            cve_info.append(f"CVE: {cve['cve_id']}")
        if 'cvss_score' in cve:
            cve_info.append(f"CVSS: {cve['cvss_score']}")
        if 'severity' in cve:
            cve_info.append(f"Severity: {cve['severity']}")
        if 'description' in cve:
            cve_info.append(f"Description: {cve['description'][:100]}...")
        
        formatted_items.append(" | ".join(cve_info))
    
    return "\n".join([f"- {item}" for item in formatted_items])


def _format_network_data(network_data: List[Dict]) -> str:
    """Format network-specific data for prompts."""
    if not network_data:
        return "No network data available"
    
    formatted_items = []
    for item in network_data:
        net_info = []
        if 'ip' in item:
            net_info.append(f"IP: {item['ip']}")
        if 'ports' in item:
            net_info.append(f"Ports: {item['ports']}")
        if 'services' in item:
            net_info.append(f"Services: {item['services']}")
        if 'country' in item:
            net_info.append(f"Country: {item['country']}")
        if 'org' in item:
            net_info.append(f"Org: {item['org']}")
        
        formatted_items.append(" | ".join(net_info))
    
    return "\n".join([f"- {item}" for item in formatted_items])


def get_custom_prompt(template_name: str, entity: str, data: Any, **kwargs) -> str:
    """
    Generate a custom prompt based on template name and parameters.
    
    Args:
        template_name: Name of the prompt template to use
        entity: Target entity
        data: Analysis data
        **kwargs: Additional template parameters
    
    Returns:
        Formatted prompt string
    """
    
    templates = {
        'summary': get_summary_prompt,
        'threat': get_threat_analysis_prompt,
        'vulnerability': get_vulnerability_analysis_prompt,
        'network': get_network_analysis_prompt,
        'correlation': get_correlation_analysis_prompt,
        'incident': get_incident_response_prompt
    }
    
    if template_name not in templates:
        raise ValueError(f"Unknown template: {template_name}")
    
    template_func = templates[template_name]
    
    # Handle different function signatures
    if template_name == 'summary':
        risk_assessment = kwargs.get('risk_assessment', {})
        return template_func(entity, data, risk_assessment)
    elif template_name in ['threat', 'network']:
        return template_func(entity, data)
    elif template_name == 'vulnerability':
        asset_data = kwargs.get('asset_data', [])
        return template_func(entity, data, asset_data)
    elif template_name == 'correlation':
        return template_func(entity, data)
    elif template_name == 'incident':
        urgency = kwargs.get('urgency', 'medium')
        return template_func(entity, data, urgency)
    
    return template_func(entity, data)