auto_sigma_rule_generator/backend/setup_ollama_with_sigma.py

#!/usr/bin/env python3
"""
Enhanced Ollama setup script that includes SIGMA model creation.
This integrates the functionality from fix_sigma_model.sh into the Docker container.
"""
import os
import json
import time
import requests
import subprocess
import sys
import tempfile
from typing import Dict, List, Optional

OLLAMA_BASE_URL = os.getenv('OLLAMA_BASE_URL', 'http://ollama:11434')
DEFAULT_MODEL = os.getenv('LLM_MODEL', 'llama3.2')
SIGMA_MODEL_NAME = 'sigma-llama'

def log(message: str, level: str = "INFO"):
    """Log message with timestamp"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {level}: {message}")

def wait_for_ollama(max_retries: int = 30, delay: int = 5) -> bool:
    """Wait for Ollama service to be ready"""
    log("Waiting for Ollama service to be ready...")
    
    for attempt in range(max_retries):
        try:
            response = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=10)
            if response.status_code == 200:
                log("✅ Ollama service is ready!")
                return True
        except requests.exceptions.RequestException as e:
            log(f"Attempt {attempt + 1}/{max_retries}: Ollama not ready yet ({e})", "DEBUG")
        
        if attempt < max_retries - 1:
            time.sleep(delay)
    
    log("❌ Ollama service failed to become ready", "ERROR")
    return False

def get_available_models() -> List[str]:
    """Get list of available models"""
    try:
        response = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=10)
        if response.status_code == 200:
            data = response.json()
            models = [model.get('name', '') for model in data.get('models', [])]
            log(f"Available models: {models}")
            return models
        else:
            log(f"Failed to get models: HTTP {response.status_code}", "ERROR")
            return []
    except Exception as e:
        log(f"Error getting models: {e}", "ERROR")
        return []

def pull_model(model_name: str) -> bool:
    """Pull a model if not available"""
    log(f"Pulling model: {model_name}")
    
    try:
        payload = {"name": model_name}
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/pull",
            json=payload,
            timeout=600,  # 10 minutes timeout
            stream=True
        )
        
        if response.status_code == 200:
            # Stream and log progress
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        status = data.get('status', '')
                        if status:
                            log(f"Pull progress: {status}", "DEBUG")
                        if data.get('error'):
                            log(f"Pull error: {data.get('error')}", "ERROR")
                            return False
                    except json.JSONDecodeError:
                        continue
            
            log(f"✅ Successfully pulled model: {model_name}")
            return True
        else:
            log(f"❌ Failed to pull model {model_name}: HTTP {response.status_code}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error pulling model {model_name}: {e}", "ERROR")
        return False

def create_sigma_model() -> bool:
    """Create the sigma-llama model with specialized SIGMA generation configuration"""
    log("🔄 Creating sigma-llama model...")
    
    # First, remove any existing sigma-llama model
    try:
        response = requests.delete(f"{OLLAMA_BASE_URL}/api/delete", 
                                 json={"name": SIGMA_MODEL_NAME},
                                 timeout=30)
        if response.status_code == 200:
            log("Removed existing sigma-llama model")
    except Exception:
        pass  # Model might not exist, that's fine
    
    # Create Modelfile content without the FROM line
    modelfile_content = """TEMPLATE \"\"\"### Instruction:
Generate SIGMA rule logsource and detection sections based on the provided context.

### Input:
{{ .Prompt }}

### Response:
\"\"\"

PARAMETER temperature 0.1
PARAMETER top_p 0.9
PARAMETER stop "### Instruction:"
PARAMETER stop "### Response:"
PARAMETER num_ctx 4096

SYSTEM \"\"\"You are a cybersecurity expert specializing in SIGMA rule creation. Generate valid SIGMA rules in YAML format based on the provided CVE and exploit information. Output ONLY valid YAML starting with 'title:' and ending with the last YAML line.\"\"\"
"""
    
    try:
        # Create the model using the API with 'from' parameter
        payload = {
            "name": SIGMA_MODEL_NAME,
            "from": f"{DEFAULT_MODEL}:latest",
            "modelfile": modelfile_content,
            "stream": False
        }
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/create",
            json=payload,
            timeout=300,  # 5 minutes timeout
            stream=True
        )
        
        if response.status_code == 200:
            # Stream and log progress
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        status = data.get('status', '')
                        if status:
                            log(f"Model creation: {status}", "DEBUG")
                        if data.get('error'):
                            log(f"Model creation error: {data.get('error')}", "ERROR")
                            return False
                    except json.JSONDecodeError:
                        continue
            
            # Verify the model was created
            models = get_available_models()
            if any(SIGMA_MODEL_NAME in model for model in models):
                log("✅ sigma-llama model created successfully!")
                return True
            else:
                log("❌ sigma-llama model not found after creation", "ERROR")
                return False
        else:
            log(f"❌ Failed to create sigma-llama model: HTTP {response.status_code}", "ERROR")
            try:
                error_data = response.json()
                log(f"Error details: {error_data}", "ERROR")
            except:
                log(f"Error response: {response.text}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error creating sigma-llama model: {e}", "ERROR")
        return False

def test_sigma_model() -> bool:
    """Test the sigma-llama model"""
    log("🔄 Testing sigma-llama model...")
    
    try:
        test_payload = {
            "model": SIGMA_MODEL_NAME,
            "prompt": "Title: Test PowerShell Rule",
            "stream": False
        }
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/generate",
            json=test_payload,
            timeout=60
        )
        
        if response.status_code == 200:
            data = response.json()
            test_response = data.get('response', '')[:100]  # First 100 chars
            log(f"✅ Model test successful! Response: {test_response}...")
            return True
        else:
            log(f"❌ Model test failed: HTTP {response.status_code}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error testing model: {e}", "ERROR")
        return False

def main():
    """Main setup function"""
    log("🚀 Starting enhanced Ollama setup with SIGMA model creation...")
    
    # Step 1: Wait for Ollama to be ready
    if not wait_for_ollama():
        log("❌ Setup failed: Ollama service not available", "ERROR")
        sys.exit(1)
    
    # Step 2: Check current models
    models = get_available_models()
    log(f"Current models: {models}")
    
    # Step 3: Pull default model if needed
    if not any(DEFAULT_MODEL in model for model in models):
        log(f"Default model {DEFAULT_MODEL} not found, pulling...")
        if not pull_model(DEFAULT_MODEL):
            log(f"❌ Setup failed: Could not pull {DEFAULT_MODEL}", "ERROR")
            sys.exit(1)
    else:
        log(f"✅ Default model {DEFAULT_MODEL} already available")
    
    # Step 4: Create SIGMA model
    if not create_sigma_model():
        log("❌ Setup failed: Could not create sigma-llama model", "ERROR")
        sys.exit(1)
    
    # Step 5: Test SIGMA model
    if not test_sigma_model():
        log("⚠️  Setup warning: sigma-llama model test failed", "WARN")
        # Don't exit here, the model might still work
    
    # Step 6: Final verification
    final_models = get_available_models()
    log(f"Final models available: {final_models}")
    
    if any(SIGMA_MODEL_NAME in model for model in final_models):
        log("🎉 Setup complete! sigma-llama model is ready for use.")
        sys.exit(0)
    else:
        log("❌ Setup failed: sigma-llama model not available after setup", "ERROR")
        sys.exit(1)

if __name__ == "__main__":
    main()