predictor_api.py

#!/usr/bin/env python3
"""
Baskerville Solo - Predictor API
Simple REST API for receiving web logs and simulating predictions to PostgreSQL
"""

import logging
import os
import time
from datetime import datetime
from typing import List, Dict, Any
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import psycopg2
from psycopg2.extras import RealDictCursor
import uvicorn

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# Pydantic models for request validation
class WebLogRequest(BaseModel):
    """Single web log entry"""
    timestamp: str
    ip: str
    host: str
    path: str
    method: str = "GET"
    status: int = 200
    user_agent: str = ""
    country: str = ""
    session_id: str = ""


class WebLogBatch(BaseModel):
    """Batch of web logs"""
    logs: List[WebLogRequest]


# Database connection
class DatabaseManager:
    def __init__(self):
        self.conn = None
        self.connect()

    def connect(self):
        """Connect to PostgreSQL"""
        try:
            self.conn = psycopg2.connect(
                host=os.getenv("POSTGRES_HOST", "postgres"),
                port=os.getenv("POSTGRES_PORT", "5432"),
                database=os.getenv("POSTGRES_DB", "baskerville"),
                user=os.getenv("POSTGRES_USER", "baskerville"),
                password=os.getenv("POSTGRES_PASSWORD", "baskerville123"),
            )
            self.conn.autocommit = True
            logger.info("Connected to PostgreSQL")
        except Exception as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}")
            raise

    def ensure_tables(self):
        """Create tables if they don't exist"""
        with self.conn.cursor() as cur:
            # Table for received logs
            cur.execute("""
                CREATE TABLE IF NOT EXISTS raw_logs (
                    id SERIAL PRIMARY KEY,
                    timestamp TIMESTAMP NOT NULL,
                    ip VARCHAR(45) NOT NULL,
                    host VARCHAR(255) NOT NULL,
                    path TEXT NOT NULL,
                    method VARCHAR(10),
                    status INTEGER,
                    user_agent TEXT,
                    country VARCHAR(10),
                    session_id VARCHAR(255),
                    received_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Table for predictions (simulation)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS predictions (
                    id SERIAL PRIMARY KEY,
                    session_id VARCHAR(255) NOT NULL,
                    ip VARCHAR(45) NOT NULL,
                    host VARCHAR(255) NOT NULL,
                    prediction VARCHAR(50) NOT NULL,
                    score_if FLOAT,
                    score_ae FLOAT,
                    is_anomaly BOOLEAN,
                    command VARCHAR(50),
                    metadata JSONB,
                    predicted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Index for faster queries
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_predictions_ip
                ON predictions(ip)
            """)
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_predictions_host
                ON predictions(host)
            """)

            # Migration: Add command column if it doesn't exist
            cur.execute("""
                DO $$
                BEGIN
                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns
                        WHERE table_name = 'predictions' AND column_name = 'command'
                    ) THEN
                        ALTER TABLE predictions ADD COLUMN command VARCHAR(50);
                    END IF;
                END $$;
            """)

            logger.info("Database tables ensured")

    def insert_log(self, log: WebLogRequest):
        """Insert a single log entry"""
        with self.conn.cursor() as cur:
            cur.execute("""
                INSERT INTO raw_logs
                (timestamp, ip, host, path, method, status, user_agent, country, session_id)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                log.timestamp,
                log.ip,
                log.host,
                log.path,
                log.method,
                log.status,
                log.user_agent,
                log.country,
                log.session_id
            ))

    def insert_prediction(self, session_id: str, ip: str, host: str,
                         prediction: str, score_if: float, score_ae: float,
                         is_anomaly: bool, command: str, metadata: Dict[str, Any]):
        """Insert a prediction (simulation)"""
        import json
        with self.conn.cursor() as cur:
            cur.execute("""
                INSERT INTO predictions
                (session_id, ip, host, prediction, score_if, score_ae, is_anomaly, command, metadata)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                session_id,
                ip,
                host,
                prediction,
                score_if,
                score_ae,
                is_anomaly,
                command,
                json.dumps(metadata)
            ))

    def get_recent_predictions(self, limit: int = 10):
        """Get recent predictions"""
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("""
                SELECT * FROM predictions
                ORDER BY predicted_at DESC
                LIMIT %s
            """, (limit,))
            return cur.fetchall()

    def close(self):
        """Close connection"""
        if self.conn:
            self.conn.close()
            logger.info("Database connection closed")


# Global database manager
db_manager = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown events"""
    global db_manager

    # Startup
    logger.info("Starting Baskerville Solo Predictor API...")

    # Wait for PostgreSQL to be ready
    max_retries = 30
    for i in range(max_retries):
        try:
            db_manager = DatabaseManager()
            db_manager.ensure_tables()
            break
        except Exception as e:
            if i < max_retries - 1:
                logger.warning(f"Waiting for PostgreSQL... ({i+1}/{max_retries})")
                time.sleep(2)
            else:
                raise

    logger.info("API is ready to receive logs")

    yield

    # Shutdown
    logger.info("Shutting down...")
    if db_manager:
        db_manager.close()


# Create FastAPI app
app = FastAPI(
    title="Baskerville Solo - Predictor API",
    description="REST API for receiving web logs and making predictions",
    version="0.1.0",
    lifespan=lifespan
)


@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "running",
        "service": "baskerville-solo-predictor",
        "version": "0.1.0"
    }


@app.get("/health")
async def health_check():
    """Detailed health check"""
    try:
        # Check database connection
        with db_manager.conn.cursor() as cur:
            cur.execute("SELECT 1")

        return {
            "status": "healthy",
            "database": "connected",
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"Unhealthy: {str(e)}")


@app.post("/api/v1/logs")
async def receive_log(log: WebLogRequest):
    """
    Receive a single web log entry

    This simulates:
    1. Storing the log
    2. Creating a prediction
    3. Storing the prediction in PostgreSQL
    """
    try:
        # Store the log
        db_manager.insert_log(log)
        logger.info(f"Received log: {log.ip} -> {log.host}{log.path}")

        # Simulate prediction
        # In real implementation, this would:
        # - Group logs into sessions
        # - Extract features
        # - Run ML models
        # - Make decision

        import random
        score_if = random.uniform(-1.0, 1.0)  # Isolation Forest score
        score_ae = random.uniform(0.0, 2.0)   # AutoEncoder score
        is_anomaly = score_if < 0 or score_ae > 1.0

        prediction = "anomaly" if is_anomaly else "normal"

        # Store prediction only if anomaly detected
        if is_anomaly:
            # Randomly select WAF command
            command = random.choice(["challenge_ip", "block_ip"])

            db_manager.insert_prediction(
                session_id=log.session_id or f"sim_{log.ip}_{int(time.time())}",
                ip=log.ip,
                host=log.host,
                prediction=prediction,
                score_if=score_if,
                score_ae=score_ae,
                is_anomaly=is_anomaly,
                command=command,
                metadata={
                    "user_agent": log.user_agent,
                    "country": log.country,
                    "path": log.path,
                    "status": log.status
                }
            )

        return {
            "status": "success",
            "message": "Log received and processed",
            "prediction": {
                "result": prediction,
                "score_if": round(score_if, 3),
                "score_ae": round(score_ae, 3),
                "is_anomaly": is_anomaly
            }
        }

    except Exception as e:
        logger.error(f"Error processing log: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/logs/batch")
async def receive_logs_batch(batch: WebLogBatch):
    """
    Receive a batch of web logs

    More efficient for high-volume scenarios
    """
    try:
        processed = 0
        predictions_made = 0

        for log in batch.logs:
            # Store log
            db_manager.insert_log(log)
            processed += 1

            # Simulate prediction (randomly for demo)
            import random
            if random.random() > 0.5:  # 50% chance to make prediction
                score_if = random.uniform(-1.0, 1.0)
                score_ae = random.uniform(0.0, 2.0)
                is_anomaly = score_if < 0 or score_ae > 1.0

                # Store prediction only if anomaly detected
                if is_anomaly:
                    # Randomly select WAF command
                    command = random.choice(["challenge_ip", "block_ip"])

                    db_manager.insert_prediction(
                        session_id=log.session_id or f"sim_{log.ip}_{int(time.time())}",
                        ip=log.ip,
                        host=log.host,
                        prediction="anomaly",
                        score_if=score_if,
                        score_ae=score_ae,
                        is_anomaly=is_anomaly,
                        command=command,
                        metadata={"batch_processed": True}
                    )
                    predictions_made += 1

        logger.info(f"Processed batch: {processed} logs, {predictions_made} predictions")

        return {
            "status": "success",
            "logs_received": processed,
            "predictions_made": predictions_made
        }

    except Exception as e:
        logger.error(f"Error processing batch: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/v1/predictions")
async def get_predictions(limit: int = 10):
    """
    Get recent predictions
    """
    try:
        predictions = db_manager.get_recent_predictions(limit)
        return {
            "status": "success",
            "count": len(predictions),
            "predictions": predictions
        }
    except Exception as e:
        logger.error(f"Error fetching predictions: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/v1/stats")
async def get_stats():
    """
    Get basic statistics
    """
    try:
        with db_manager.conn.cursor(cursor_factory=RealDictCursor) as cur:
            # Count logs
            cur.execute("SELECT COUNT(*) as total FROM raw_logs")
            log_count = cur.fetchone()['total']

            # Count predictions
            cur.execute("SELECT COUNT(*) as total FROM predictions")
            prediction_count = cur.fetchone()['total']

            # Count anomalies
            cur.execute("SELECT COUNT(*) as total FROM predictions WHERE is_anomaly = true")
            anomaly_count = cur.fetchone()['total']

            return {
                "status": "success",
                "stats": {
                    "total_logs": log_count,
                    "total_predictions": prediction_count,
                    "total_anomalies": anomaly_count,
                    "anomaly_rate": round(anomaly_count / prediction_count * 100, 2) if prediction_count > 0 else 0
                }
            }
    except Exception as e:
        logger.error(f"Error fetching stats: {e}")
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=int(os.getenv("API_PORT", "8000")),
        log_level="info"
    )