src/app.py

import asyncio
import os
import logging
from typing import List, Dict, Any, Optional
from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from dotenv import load_dotenv

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

load_dotenv()

from connectors.intelligencex import search_intelx
from connectors.enisa_cve import fetch_enisa_cve
#from connectors.maltego import fetch_from_maltego
#from connectors.shodan import search_shodan
from connectors.virustotal import search_virustotal

from core.normalize import normalize_intelx_data, normalize_shodan_data, normalize_virustotal_data
from core.correlate import correlate_entities
from core.risk_engine import assess_risk
from llm_agent.llm_chatGPT import load_llm
from llm_agent.prompt_templates import get_summary_prompt

app = FastAPI(
    title="OSINT Early Warning System",
    description="An intelligence gathering and analysis API",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Consider restricting this in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files for frontend
frontend_path = os.path.join(os.path.dirname(__file__), "../osint-early-warning/frontend/build")
if os.path.exists(frontend_path):
    app.mount(
        "/",
        StaticFiles(directory=frontend_path, html=True),
        name="frontend",
    )
else:
    logger.warning(f"Frontend directory not found: {frontend_path}")

class SummaryResponse(BaseModel):
    summary: str
    risk_level: str
    confidence: float
    sources_used: List[str]
    timestamp: str

class ErrorResponse(BaseModel):
    error: str
    details: Optional[str] = None

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "version": "1.0.0"}

@app.get("/api/osint", response_model=SummaryResponse, responses={400: {"model": ErrorResponse}})
async def osint_summary(entity: str = Query(..., min_length=1, max_length=255)):
    """
    Perform OSINT analysis on a given entity
    
    Args:
        entity: The target entity to analyze (IP, domain, hash, etc.)
    
    Returns:
        Comprehensive analysis summary with risk assessment
    """
    try:
        logger.info(f"Starting OSINT analysis for entity: {entity}")
        
        # Initialize data containers
        all_normalized_data = []
        sources_used = []
        
        # Gather data from various sources concurrently
        tasks = []
        
        # IntelligenceX
        tasks.append(("intelx", search_intelx(entity)))
        
        # ENISA CVE
        tasks.append(("enisa_cve", fetch_enisa_cve(entity)))
        
        # VirusTotal
        tasks.append(("virustotal", search_virustotal(entity)))
        
        # Uncomment when ready to use
        # tasks.append(("shodan", search_shodan(entity)))
        # tasks.append(("maltego", fetch_from_maltego(entity)))
        
        # Execute all tasks concurrently
        results = {}
        for source_name, task in tasks:
            try:
                result = await task
                results[source_name] = result
                if result:  # Only add to sources if we got data
                    sources_used.append(source_name)
                logger.info(f"Successfully gathered data from {source_name}")
            except Exception as e:
                logger.error(f"Error gathering data from {source_name}: {str(e)}")
                results[source_name] = None
        
        # Normalize data from successful sources
        if results.get("intelx"):
            try:
                intelx_normalized = normalize_intelx_data(results["intelx"])
                all_normalized_data.extend(intelx_normalized)
            except Exception as e:
                logger.error(f"Error normalizing IntelX data: {str(e)}")
        
        if results.get("virustotal"):
            try:
                vt_normalized = normalize_virustotal_data(results["virustotal"])
                all_normalized_data.extend(vt_normalized)
            except Exception as e:
                logger.error(f"Error normalizing VirusTotal data: {str(e)}")
        
        # Uncomment when Shodan is ready
        # if results.get("shodan"):
        #     try:
        #         shodan_normalized = normalize_shodan_data(results["shodan"])
        #         all_normalized_data.extend(shodan_normalized)
        #     except Exception as e:
        #         logger.error(f"Error normalizing Shodan data: {str(e)}")
        
        # Check if we have any data to work with
        if not all_normalized_data and not results.get("enisa_cve"):
            raise HTTPException(
                status_code=404, 
                detail="No data found for the specified entity"
            )
        
        # Perform correlation analysis
        try:
            correlated_data = correlate_entities(results.get("enisa_cve", []), all_normalized_data)
            logger.info("Successfully correlated entity data")
        except Exception as e:
            logger.error(f"Error correlating data: {str(e)}")
            correlated_data = all_normalized_data
        
        # Assess risk
        try:
            risk_assessment = assess_risk(correlated_data)
            logger.info(f"Risk assessment completed: {risk_assessment}")
        except Exception as e:
            logger.error(f"Error assessing risk: {str(e)}")
            risk_assessment = {"level": "unknown", "confidence": 0.0}
        
        # Generate LLM summary
        try:
            llm = load_llm()
            prompt = get_summary_prompt(entity, correlated_data, risk_assessment)
            summary = await llm(prompt) if asyncio.iscoroutinefunction(llm) else llm(prompt)
            logger.info("Successfully generated LLM summary")
        except Exception as e:
            logger.error(f"Error generating LLM summary: {str(e)}")
            summary = f"Analysis completed for {entity}. Manual review recommended due to processing error."
        
        # Prepare response
        from datetime import datetime
        response = SummaryResponse(
            summary=summary,
            risk_level=risk_assessment.get("level", "unknown"),
            confidence=risk_assessment.get("confidence", 0.0),
            sources_used=sources_used,
            timestamp=datetime.utcnow().isoformat()
        )
        
        logger.info(f"OSINT analysis completed successfully for entity: {entity}")
        return response
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in OSINT analysis: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail="Internal server error during analysis"
        )

@app.get("/api/sources")
async def get_available_sources():
    """Get list of available OSINT sources"""
    sources = [
        {"name": "IntelligenceX", "status": "active", "type": "threat_intelligence"},
        {"name": "ENISA CVE", "status": "active", "type": "vulnerability"},
        {"name": "VirusTotal", "status": "active", "type": "malware_analysis"},
        {"name": "Shodan", "status": "inactive", "type": "network_scanning"},
        {"name": "Maltego", "status": "inactive", "type": "graph_analysis"},
    ]
    return {"sources": sources}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")