49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from urllib import request
|
|
|
|
from .models import BlockCandidate, Finding
|
|
|
|
|
|
def ollama_summary(
|
|
findings: list[Finding],
|
|
candidates: list[BlockCandidate],
|
|
model: str | None = None,
|
|
*,
|
|
analysis: dict[str, object] | None = None,
|
|
timeout: int | None = None,
|
|
) -> str:
|
|
host = os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
|
|
selected_model = model or os.getenv("OLLAMA_MODEL", "llama3.3")
|
|
selected_timeout = timeout or int(os.getenv("OLLAMA_TIMEOUT", "180"))
|
|
prompt = {
|
|
"analysis": analysis or {},
|
|
"findings": [finding.__dict__ for finding in findings],
|
|
"block_candidates": [
|
|
{"src_ip": candidate.src_ip, "score": candidate.score, "reasons": candidate.reasons}
|
|
for candidate in candidates[:20]
|
|
],
|
|
}
|
|
body = json.dumps(
|
|
{
|
|
"model": selected_model,
|
|
"stream": False,
|
|
"options": {
|
|
"num_predict": 350,
|
|
"temperature": 0.2,
|
|
},
|
|
"prompt": (
|
|
"You are a local FortiGate security analyst. Summarize these policy findings "
|
|
"and log diagnostics. Be concise. Include risk, likely cause, and next action. "
|
|
"Do not recommend blocking private/internal client IPs unless the data explicitly proves compromise. "
|
|
f"Data: {json.dumps(prompt)}"
|
|
),
|
|
}
|
|
).encode("utf-8")
|
|
req = request.Request(f"{host}/api/generate", data=body, method="POST", headers={"Content-Type": "application/json"})
|
|
with request.urlopen(req, timeout=selected_timeout) as response:
|
|
data = json.loads(response.read().decode("utf-8"))
|
|
return str(data.get("response", "")).strip()
|