api_gateway.py

import time
import logging
import asyncio
import re
from fastapi import FastAPI, HTTPException, Request, Depends, Header
from pydantic import BaseModel, Field
from typing import Dict
from collections import defaultdict

# ==========================================
# 🔧 Logging
# ==========================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s"
)
logger = logging.getLogger("AegisAPI")

# ==========================================
# ⚙️ App
# ==========================================
app = FastAPI(
    title="Aegis AI Enterprise API v3",
    version="3.0.0"
)

# ==========================================
# 🔐 API Keys
# ==========================================
API_KEYS = {"test-key-123"}

async def verify_api_key(x_api_key: str = Header(...)):
    if x_api_key not in API_KEYS:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    return x_api_key

# ==========================================
# 🧠 Models
# ==========================================
class ScanRequest(BaseModel):
    text: str = Field(..., min_length=3, max_length=5000)

class ScanResponse(BaseModel):
    status: str
    risk_level: str
    score: float
    action: str
    threat_score: float
    timestamp: float

# ==========================================
# 📊 Rate Limiting (IP + API KEY)
# ==========================================
RATE_LIMIT = 20
WINDOW = 60
client_requests: Dict[str, list] = defaultdict(list)

def check_rate_limit(identifier: str):
    now = time.time()
    client_requests[identifier] = [
        t for t in client_requests[identifier] if now - t < WINDOW
    ]

    if len(client_requests[identifier]) >= RATE_LIMIT:
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    client_requests[identifier].append(now)

# ==========================================
# 🧠 Active Defender (FIXED)
# ==========================================
class ActiveDefender:
    def __init__(self):
        self.memory = defaultdict(lambda: {"count": 0, "last": 0})

    def decay(self, client_id):
        now = time.time()
        data = self.memory[client_id]

        if now - data["last"] > 120:
            data["count"] = max(0, data["count"] - 1)

    def register(self, client_id):
        data = self.memory[client_id]
        data["count"] = min(data["count"] + 1, 20)
        data["last"] = time.time()

    def adjust(self, client_id, score):
        if score <= 0:
            return 0.0

        attempts = self.memory[client_id]["count"]
        return min(score + attempts * 0.03, 1.0)

    async def get_action(self, client_id, score):
        self.decay(client_id)

        if score >= 0.5:
            self.register(client_id)

        adjusted = self.adjust(client_id, score)

        if adjusted >= 0.95:
            return "BLOCK_AND_REPORT"

        if adjusted >= 0.80:
            delay = min(0.5 * self.memory[client_id]["count"], 3)
            await asyncio.sleep(delay)
            return "THROTTLE"

        if adjusted >= 0.60:
            return "CHALLENGE"

        return "PASS"

defender = ActiveDefender()

# ==========================================
# 🔍 Detection Engine (Regex-based)
# ==========================================
PATTERNS = {
    "SQL": re.compile(r"(union\s+select|select\s+.*\s+from|drop\s+table)", re.I),
    "XSS": re.compile(r"(<script|javascript:|onerror=)", re.I),
    "PROMPT": re.compile(r"(ignore\s+all\s+previous|system\s+prompt)", re.I),
    "LINK": re.compile(r"https?://", re.I),
}

def analyze_text(text: str):
    normalized = text.lower()
    findings = []
    score = 0.0

    for name, pattern in PATTERNS.items():
        if pattern.search(normalized):
            findings.append(name)

            if name == "SQL":
                score += 0.9
            elif name == "XSS":
                score += 0.85
            elif name == "PROMPT":
                score += 0.8
            elif name == "LINK":
                score += 0.4

    score = min(score, 1.0)

    # 🧠 Risk level ديناميكي
    if score >= 0.9:
        risk = "CRITICAL"
    elif score >= 0.75:
        risk = "HIGH"
    elif score >= 0.5:
        risk = "MEDIUM"
    elif score > 0:
        risk = "LOW"
    else:
        risk = "SAFE"

    threat_score = min(score + len(findings) * 0.03, 1.0)

    return risk, score, threat_score, findings

# ==========================================
# 🛡️ Middleware
# ==========================================
@app.middleware("http")
async def security_layer(request: Request, call_next):
    ip = request.headers.get("X-Forwarded-For", request.client.host)
    api_key = request.headers.get("X-API-Key", "unknown")

    identifier = f"{ip}:{api_key}"
    check_rate_limit(identifier)

    response = await call_next(request)

    response.headers["X-Frame-Options"] = "DENY"
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Server"] = "AegisAI"

    return response

# ==========================================
# ❤️ Health
# ==========================================
@app.get("/health")
async def health():
    return {"status": "ok", "time": time.time()}

# ==========================================
# 🚀 Scan Endpoint
# ==========================================
@app.post("/v3/scan", response_model=ScanResponse)
async def scan(
    request: ScanRequest,
    http_request: Request,
    api_key: str = Depends(verify_api_key)
):
    try:
        ip = http_request.headers.get("X-Forwarded-For", http_request.client.host)

        logger.info({
            "event": "scan",
            "ip": ip,
            "preview": request.text[:40]
        })

        risk, score, threat_score, findings = analyze_text(request.text)

        action = await defender.get_action(ip, score)

        return ScanResponse(
            status="success",
            risk_level=risk,
            score=round(score, 2),
            threat_score=round(threat_score, 2),
            action=action,
            timestamp=time.time()
        )

    except Exception:
        logger.exception("Scan failed")
        raise HTTPException(status_code=500, detail="AI Engine Failure")

# ==========================================
# 📌 Root
# ==========================================
@app.get("/")
async def root():
    return {
        "name": "Aegis AI Enterprise",
        "version": "3.0",
        "endpoints": ["/v3/scan", "/health"]
    }