auto_sigma_rule_generator/backend/routers/llm_operations.py

from typing import Dict, Any
from fastapi import APIRouter, HTTPException, Depends
from sqlalchemy.orm import Session
from pydantic import BaseModel

from config.database import get_db
from models import CVE, SigmaRule

router = APIRouter(prefix="/api", tags=["llm-operations"])


class LLMRuleRequest(BaseModel):
    cve_id: str
    poc_content: str = ""


class LLMSwitchRequest(BaseModel):
    provider: str
    model: str = ""


@router.post("/llm-enhanced-rules")
async def generate_llm_enhanced_rules(request: LLMRuleRequest, db: Session = Depends(get_db)):
    """Generate SIGMA rules using LLM AI analysis"""
    try:
        from enhanced_sigma_generator import EnhancedSigmaGenerator
        
        # Get CVE
        cve = db.query(CVE).filter(CVE.cve_id == request.cve_id).first()
        if not cve:
            raise HTTPException(status_code=404, detail="CVE not found")
        
        # Generate enhanced rule using LLM
        generator = EnhancedSigmaGenerator(db)
        result = await generator.generate_enhanced_rule(cve, use_llm=True)
        
        if result.get('success'):
            return {
                "success": True,
                "message": f"Generated LLM-enhanced rule for {request.cve_id}",
                "rule_id": result.get('rule_id'),
                "generation_method": "llm_enhanced"
            }
        else:
            return {
                "success": False,
                "error": result.get('error', 'Unknown error'),
                "cve_id": request.cve_id
            }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating LLM-enhanced rule: {str(e)}")


@router.get("/llm-status")
async def get_llm_status():
    """Check LLM API availability and configuration for all providers"""
    try:
        from llm_client import LLMClient
        
        # Test all providers
        providers_status = {}
        
        # Test Ollama
        try:
            ollama_client = LLMClient(provider="ollama")
            ollama_status = await ollama_client.test_connection()
            providers_status["ollama"] = {
                "available": ollama_status.get("available", False),
                "models": ollama_status.get("models", []),
                "current_model": ollama_status.get("current_model"),
                "base_url": ollama_status.get("base_url")
            }
        except Exception as e:
            providers_status["ollama"] = {"available": False, "error": str(e)}
        
        # Test OpenAI
        try:
            openai_client = LLMClient(provider="openai")
            openai_status = await openai_client.test_connection()
            providers_status["openai"] = {
                "available": openai_status.get("available", False),
                "models": openai_status.get("models", []),
                "current_model": openai_status.get("current_model"),
                "has_api_key": openai_status.get("has_api_key", False)
            }
        except Exception as e:
            providers_status["openai"] = {"available": False, "error": str(e)}
        
        # Test Anthropic
        try:
            anthropic_client = LLMClient(provider="anthropic")
            anthropic_status = await anthropic_client.test_connection()
            providers_status["anthropic"] = {
                "available": anthropic_status.get("available", False),
                "models": anthropic_status.get("models", []),
                "current_model": anthropic_status.get("current_model"),
                "has_api_key": anthropic_status.get("has_api_key", False)
            }
        except Exception as e:
            providers_status["anthropic"] = {"available": False, "error": str(e)}
        
        # Get current configuration
        current_client = LLMClient()
        current_config = {
            "current_provider": current_client.provider,
            "current_model": current_client.model,
            "default_provider": "ollama"
        }
        
        return {
            "providers": providers_status,
            "configuration": current_config,
            "status": "operational" if any(p.get("available") for p in providers_status.values()) else "no_providers_available"
        }
    
    except Exception as e:
        return {
            "status": "error",
            "error": str(e),
            "providers": {},
            "configuration": {}
        }


@router.post("/llm-switch")
async def switch_llm_provider(request: LLMSwitchRequest):
    """Switch between LLM providers and models"""
    try:
        from llm_client import LLMClient
        
        # Test the new provider/model
        test_client = LLMClient(provider=request.provider, model=request.model)
        connection_test = await test_client.test_connection()
        
        if not connection_test.get("available"):
            raise HTTPException(
                status_code=400, 
                detail=f"Provider {request.provider} with model {request.model} is not available"
            )
        
        # Switch to new configuration (this would typically involve updating environment variables
        # or configuration files, but for now we'll just confirm the switch)
        return {
            "success": True,
            "message": f"Switched to {request.provider}" + (f" with model {request.model}" if request.model else ""),
            "provider": request.provider,
            "model": request.model or connection_test.get("current_model"),
            "available": True
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error switching LLM provider: {str(e)}")


@router.post("/ollama-pull-model")
async def pull_ollama_model(model: str = "llama3.2"):
    """Pull a model in Ollama"""
    try:
        import aiohttp
        import os
        
        ollama_url = os.getenv("OLLAMA_BASE_URL", "http://ollama:11434")
        
        async with aiohttp.ClientSession() as session:
            async with session.post(f"{ollama_url}/api/pull", json={"name": model}) as response:
                if response.status == 200:
                    return {
                        "success": True,
                        "message": f"Successfully pulled model {model}",
                        "model": model
                    }
                else:
                    raise HTTPException(status_code=500, detail=f"Failed to pull model: {response.status}")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error pulling Ollama model: {str(e)}")


@router.get("/ollama-models")
async def get_ollama_models():
    """Get available Ollama models"""
    try:
        import aiohttp
        import os
        
        ollama_url = os.getenv("OLLAMA_BASE_URL", "http://ollama:11434")
        
        async with aiohttp.ClientSession() as session:
            async with session.get(f"{ollama_url}/api/tags") as response:
                if response.status == 200:
                    data = await response.json()
                    models = [model["name"] for model in data.get("models", [])]
                    return {
                        "models": models,
                        "total_models": len(models),
                        "ollama_url": ollama_url
                    }
                else:
                    return {
                        "models": [],
                        "total_models": 0,
                        "error": f"Ollama not available (status: {response.status})"
                    }
    
    except Exception as e:
        return {
            "models": [],
            "total_models": 0,
            "error": f"Error connecting to Ollama: {str(e)}"
        }