genai_app.py

"""
genai_app.py — GenAI Shield V2 Flask Application.

Powered by Llama-Prompt-Guard-2-86M for pre-inference prompt screening.

Endpoints:
  GET  /                      → Chat interface
  GET  /genai-monitoring      → Real-time GenAI SIEM dashboard
  GET  /genai-stream          → SSE event stream
  POST /genai-chat            → Send a prompt, get monitored response
  GET  /guard-stats           → Prompt Guard model statistics

Configure via environment variables:
  GEMINI_API_KEY=...
  GENAI_PORT=5001  (default)
"""

import os
import json
import time
import queue
import threading
from flask import Flask, request, jsonify, render_template, Response, stream_with_context
from flask_cors import CORS

from gemini_adapter          import GeminiAdapter
from prompt_guard_engine     import PromptGuardEngine
from prompt_guard_text_guard import PromptGuardTextGuard
from text_monitor            import TextMonitor
from attachment_guard        import AttachmentGuard

app = Flask(__name__)
CORS(app)

# ── System prompt ─────────────────────────────────────────────────────────────
SYSTEM_PROMPT = os.getenv(
    "GENAI_SYSTEM_PROMPT",
    "You are a helpful AI assistant. Be concise, accurate, and professional."
)

# ── Initialise Prompt Guard engine ────────────────────────────────────────────
print("[GenAI Shield V2] Initialising Prompt Guard engine...")
PG_ENGINE = PromptGuardEngine().load()
GUARD     = PromptGuardTextGuard(PG_ENGINE)
print("[GenAI Shield V2] Prompt Guard ready.")

# ── Initialise LLM adapter + post-inference monitor ──────────────────────────
ADAPTER = GeminiAdapter(system_prompt=SYSTEM_PROMPT)
MONITOR = TextMonitor(ADAPTER, system_prompt=SYSTEM_PROMPT)
# ─────────────────────────────────────────────────────────────────────────────

# SSE Broadcast Queue
BROADCAST_QUEUES = []

def broadcast(event_type: str, data: dict):
    event = {
        "timestamp":          time.strftime("%H:%M:%S"),
        "type":               event_type,
        "threat_score":       data.get("threat_score", 0),
        "flags":              data.get("flags", []),
        "reason":             data.get("reason", "CLEAN"),
        "source":             data.get("source", "Web UI"),
        "prompt":             data.get("prompt", "")[:120],
        "response":           data.get("response", "")[:200],
        "latency_ms":         data.get("latency_ms", 0),
        "checks":             data.get("checks", {}),
        "model":              ADAPTER.get_model_name(),
        "prompt_guard_score": data.get("prompt_guard_score", 0),
    }
    for q in BROADCAST_QUEUES:
        q.put(event)


# ── Routes ────────────────────────────────────────────────────────────────────

@app.route("/")
def index():
    return render_template("genai.html", model=ADAPTER.get_model_name())

@app.route("/sidecar")
def sidecar_ui():
    """Sidecar streaming chat UI."""
    return render_template("sidecar.html")

@app.route("/dataflow")
def dataflow_ui():
    """Real-time data flow visualization dashboard."""
    return render_template("dataflow.html")

@app.route("/genai-monitoring")
def monitoring():
    return render_template("genai_monitoring.html", model=ADAPTER.get_model_name())

@app.route("/genai-stream")
def stream():
    def event_stream():
        q = queue.Queue()
        BROADCAST_QUEUES.append(q)
        try:
            while True:
                event = q.get()
                yield f"data: {json.dumps(event)}\n\n"
        except GeneratorExit:
            BROADCAST_QUEUES.remove(q)
    return Response(stream_with_context(event_stream()), mimetype="text/event-stream")


@app.route("/genai-chat", methods=["POST"])
def chat():
    data       = request.json
    prompt     = data.get("prompt", "").strip()
    attachment = data.get("attachment")   # { filename, content_b64 }
    source     = data.get("source", "Web UI")

    if not prompt:
        return jsonify({"error": "Empty prompt"}), 400

    # ── LAYER 0: Attachment Extraction & Screening ───────────────────────────
    attachment_text = ""
    if attachment:
        filename = attachment.get("filename", "unknown")
        b64      = attachment.get("content_b64", "")
        extracted = AttachmentGuard.extract_text(filename, b64)

        if extracted["error"]:
            return jsonify({"error": extracted["error"]}), 400

        attachment_text = extracted["text"]
        # Screen attachment text through Prompt Guard
        att_guard_result = AttachmentGuard.screen_with_guard(GUARD, filename, attachment_text)
        if att_guard_result["blocked"]:
            pg_score = att_guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0)
            broadcast("BLOCKED_PRE_INFERENCE", {
                "threat_score":       att_guard_result.get("threat_score", 100),
                "flags":              att_guard_result["flags"],
                "reason":             att_guard_result["reason"],
                "prompt":             prompt,
                "response":           f"[BLOCKED — Malicious Attachment: {filename}]",
                "source":             source,
                "latency_ms":         0,
                "checks":             att_guard_result.get("checks", {}),
                "prompt_guard_score": pg_score,
            })
            return jsonify({
                "blocked": True, "error": "ATTACHMENT_REJECTED",
                "reason": att_guard_result["reason"],
                "threat_score": att_guard_result.get("threat_score", 100),
                "flags": att_guard_result["flags"]
            }), 403

    # ── LAYER 1: Pre-Inference Guard (Prompt Guard model) ────────────────────
    guard_start  = time.time()
    guard_result = GUARD.screen(prompt)
    guard_lat    = round((time.time() - guard_start) * 1000, 2)

    pg_score = guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0)

    if guard_result["blocked"]:
        broadcast("BLOCKED_PRE_INFERENCE", {
            "threat_score":       guard_result["threat_score"],
            "flags":              guard_result["flags"],
            "reason":             guard_result["reason"],
            "prompt":             prompt,
            "response":           "[BLOCKED — LLM never called]",
            "source":             source,
            "latency_ms":         guard_lat,
            "checks":             guard_result["checks"],
            "prompt_guard_score": pg_score,
        })
        return jsonify({
            "blocked":      True,
            "response":     None,
            "error":        "PROMPT_REJECTED_BY_GUARD",
            "reason":       guard_result["reason"],
            "threat_score": guard_result["threat_score"],
            "flags":        guard_result["flags"],
            "prompt_guard_score": pg_score,
            "latency_breakdown": {
                "guard_ms": guard_lat,
                "model_ms": 0,
                "monitor_ms": 0
            }
        }), 403

    # ── LAYER 2: LLM Inference ────────────────────────────────────────────────
    try:
        model_start = time.time()
        full_prompt = prompt
        if attachment_text:
            full_prompt = (
                f"Context from attachment '{filename}':\n---\n{attachment_text}\n"
                f"---\nUser prompt: {prompt}"
            )
        response    = ADAPTER.chat(full_prompt, system_prompt=SYSTEM_PROMPT)
        model_lat   = round((time.time() - model_start) * 1000, 2)
    except Exception as e:
        return jsonify({"error": f"LLM error: {str(e)}"}), 500

    # ── LAYER 3: Post-Inference Monitor ───────────────────────────────────────
    monitor_start  = time.time()
    monitor_result = MONITOR.analyze(prompt, response, source=source)
    monitor_lat    = round((time.time() - monitor_start) * 1000, 2)

    total_lat = round(guard_lat + model_lat + monitor_lat, 2)

    # Determine final threat level
    threat_score = max(guard_result["threat_score"], monitor_result["threat_score"])
    all_flags    = guard_result["flags"] + monitor_result["flags"]

    # Broadcast to dashboard
    event_type = "SUSPICIOUS" if threat_score >= 30 else "INFERENCE"
    broadcast(event_type, {
        "threat_score":       threat_score,
        "flags":              all_flags,
        "reason":             monitor_result["reason"],
        "prompt":             prompt,
        "response":           response,
        "source":             source,
        "latency_ms":         total_lat,
        "prompt_guard_score": pg_score,
        "checks":             {
            "guard":   guard_result["checks"],
            "monitor": monitor_result["checks"],
            "breakdown": {
                "guard_ms":   guard_lat,
                "model_ms":   model_lat,
                "monitor_ms": monitor_lat
            }
        },
    })

    return jsonify({
        "blocked":      False,
        "response":     response,
        "threat_score": threat_score,
        "flags":        all_flags,
        "latency_ms":   total_lat,
        "prompt_guard_score": pg_score,
        "latency_breakdown": {
            "guard_ms":   guard_lat,
            "model_ms":   model_lat,
            "monitor_ms": monitor_lat
        },
        "model":        ADAPTER.get_model_name(),
    })


@app.route("/guard-stats")
def guard_stats():
    """Return Prompt Guard engine statistics."""
    return jsonify(PG_ENGINE.stats())


def _start_sidecar_subprocess():
    """
    Optionally launch the sidecar as a sub-process so both UIs run together.
    Controlled via LAUNCH_SIDECAR=true env var.
    """
    import subprocess, sys
    sidecar_port = int(os.getenv("SIDECAR_PORT", 5050))
    print(f"[GenAI Shield] Launching sidecar on :{sidecar_port}...")
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "sidecar.app:app",
         "--host", "0.0.0.0", "--port", str(sidecar_port),
         "--log-level", "warning"],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )
    def _pipe():
        for line in proc.stdout:
            print("[sidecar]", line.decode(errors='replace').rstrip())
    threading.Thread(target=_pipe, daemon=True).start()
    return proc


if __name__ == "__main__":
    port = int(os.getenv("GENAI_PORT", 5001))
    print(f"GenAI Shield V2 starting on port {port}")
    print(f"LLM Model: {ADAPTER.get_model_name()}")
    print(f"Guard:     Llama-Prompt-Guard-2-86M")
    print(f"Sidecar UI: http://localhost:{port}/sidecar")

    if os.getenv("LAUNCH_SIDECAR", "").lower() in ("1", "true", "yes"):
        _start_sidecar_subprocess()

    app.run(host="0.0.0.0", port=port, debug=False)