From b4ad932c9105c6f9eb98ba4b9c82eed414196fbceddaf90e0737e33c522c79bc Mon Sep 17 00:00:00 2001 From: larssand Date: Thu, 18 Jun 2026 22:17:03 +0200 Subject: [PATCH] add timoute och export av fg config --- README.md | 16 +++++++++++- src/fgai/cli.py | 37 ++++++++++++++-------------- src/fgai/llm.py | 22 ++++++++++++++--- src/fgai/policies.py | 55 +++++++++++++++++++++++++++++++++++++++++- tests/test_policies.py | 25 ++++++++++++++++++- 5 files changed, 130 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 006f686..9bab077 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,13 @@ Audit a FortiGate policy export: fgai audit-policies --config exports/fortigate.conf ``` +Or fetch policies through the FortiGate API and audit that JSON: + +```bash +fgai fetch-policies --output exports/policies.json +fgai audit-policies --config exports/policies.json --llm --llm-timeout 300 +``` + Find block candidates without changing the firewall: ```bash @@ -97,7 +104,13 @@ Optional local LLM summary through Ollama: ```bash ollama pull llama3.3 -fgai analyze-logs --logs logs/fg_syslog.jsonl --llm +fgai analyze-logs --logs logs/fg_syslog.jsonl --llm --llm-timeout 300 +``` + +For slower machines or large models: + +```bash +OLLAMA_MODEL=llama3.1 OLLAMA_TIMEOUT=300 fgai analyze-logs --logs logs/fg_syslog.jsonl --llm ``` ## FortiGate Inputs @@ -126,6 +139,7 @@ end - `FGAI_ALLOWLIST`: comma-separated IPs/CIDRs never to block. - `OLLAMA_HOST`: defaults to `http://127.0.0.1:11434`. - `OLLAMA_MODEL`: defaults to `llama3.3`. +- `OLLAMA_TIMEOUT`: Ollama request timeout in seconds, defaults to `180`. ## Safety Model diff --git a/src/fgai/cli.py b/src/fgai/cli.py index b5eb700..f692abe 100644 --- a/src/fgai/cli.py +++ b/src/fgai/cli.py @@ -24,25 +24,24 @@ def analyze_logs(args: argparse.Namespace) -> int: min_score=args.min_score, allowlist=parse_allowlist(args.allowlist), ) - _print_json( - { - "summary": summarize_events(events), - "diagnostics": { - "top_source_ips": top_field_values(events, "srcip", limit=10), - "top_services": top_field_values(events, "service", limit=10), - "top_actions": top_field_values(events, "action", limit=10), - "top_subtypes": top_field_values(events, "subtype", limit=10), - "local_in_failures": local_in_failures(events, limit=10), - }, - "block_candidates": [ - {"src_ip": candidate.src_ip, "score": candidate.score, "reasons": candidate.reasons} - for candidate in candidates - ], - } - ) + analysis = { + "summary": summarize_events(events), + "diagnostics": { + "top_source_ips": top_field_values(events, "srcip", limit=10), + "top_services": top_field_values(events, "service", limit=10), + "top_actions": top_field_values(events, "action", limit=10), + "top_subtypes": top_field_values(events, "subtype", limit=10), + "local_in_failures": local_in_failures(events, limit=10), + }, + "block_candidates": [ + {"src_ip": candidate.src_ip, "score": candidate.score, "reasons": candidate.reasons} + for candidate in candidates + ], + } + _print_json(analysis) if args.llm: print("\nLLM summary:") - print(ollama_summary([], candidates, args.model)) + print(ollama_summary([], candidates, args.model, analysis=analysis, timeout=args.llm_timeout)) return 0 @@ -51,7 +50,7 @@ def audit_policy_file(args: argparse.Namespace) -> int: _print_json([finding.__dict__ for finding in findings]) if args.llm: print("\nLLM summary:") - print(ollama_summary(findings, [], args.model)) + print(ollama_summary(findings, [], args.model, timeout=args.llm_timeout)) return 0 @@ -117,12 +116,14 @@ def build_parser() -> argparse.ArgumentParser: logs.add_argument("--allowlist", default=None, help="Comma-separated IPs/CIDRs never to block") logs.add_argument("--llm", action="store_true", help="Ask local Ollama to summarize results") logs.add_argument("--model", default=None, help="Ollama model name") + logs.add_argument("--llm-timeout", type=int, default=None, help="Ollama request timeout in seconds") logs.set_defaults(func=analyze_logs) policies = subparsers.add_parser("audit-policies", help="Audit FortiOS firewall policy config") policies.add_argument("--config", required=True, help="Path to FortiOS config backup") policies.add_argument("--llm", action="store_true", help="Ask local Ollama to summarize results") policies.add_argument("--model", default=None, help="Ollama model name") + policies.add_argument("--llm-timeout", type=int, default=None, help="Ollama request timeout in seconds") policies.set_defaults(func=audit_policy_file) blocks = subparsers.add_parser("suggest-blocks", help="Suggest or execute guarded source IP blocks") diff --git a/src/fgai/llm.py b/src/fgai/llm.py index 5530280..7899db5 100644 --- a/src/fgai/llm.py +++ b/src/fgai/llm.py @@ -7,28 +7,42 @@ from urllib import request from .models import BlockCandidate, Finding -def ollama_summary(findings: list[Finding], candidates: list[BlockCandidate], model: str | None = None) -> str: +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 + 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 UTM block candidates. Be concise, include risk, likely cause, and next action. " + "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=60) as response: + with request.urlopen(req, timeout=selected_timeout) as response: data = json.loads(response.read().decode("utf-8")) return str(data.get("response", "")).strip() diff --git a/src/fgai/policies.py b/src/fgai/policies.py index 3a48d5d..6524561 100644 --- a/src/fgai/policies.py +++ b/src/fgai/policies.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import shlex from pathlib import Path @@ -56,8 +57,60 @@ def parse_policy_config(text: str) -> list[PolicyRule]: return policies +def _normalize_policy_value(value: object) -> list[str]: + if value is None: + return [] + if isinstance(value, list): + normalized: list[str] = [] + for item in value: + if isinstance(item, dict): + if "name" in item: + normalized.append(str(item["name"])) + elif "q_origin_key" in item: + normalized.append(str(item["q_origin_key"])) + else: + normalized.append(json.dumps(item, sort_keys=True)) + else: + normalized.append(str(item)) + return normalized + if isinstance(value, dict): + if "name" in value: + return [str(value["name"])] + if "q_origin_key" in value: + return [str(value["q_origin_key"])] + return [json.dumps(value, sort_keys=True)] + return [str(value)] + + +def parse_policy_json(text: str) -> list[PolicyRule]: + data = json.loads(text) + raw_policies = data.get("results", data) if isinstance(data, dict) else data + if not isinstance(raw_policies, list): + raise ValueError("policy JSON must be a list or contain a results list") + + policies: list[PolicyRule] = [] + for index, raw_policy in enumerate(raw_policies, start=1): + if not isinstance(raw_policy, dict): + continue + policy_id = str(raw_policy.get("policyid", raw_policy.get("policy_id", raw_policy.get("id", index)))) + settings = {str(key): _normalize_policy_value(value) for key, value in raw_policy.items()} + policies.append(PolicyRule(policy_id, settings)) + return policies + + def read_policies(path: str | Path) -> list[PolicyRule]: - return parse_policy_config(Path(path).read_text(encoding="utf-8", errors="replace")) + policy_path = Path(path) + if not policy_path.exists(): + raise FileNotFoundError( + f"{policy_path} does not exist. Fetch policies first with: " + f"fgai fetch-policies --output exports/policies.json" + ) + + text = policy_path.read_text(encoding="utf-8", errors="replace") + stripped = text.lstrip() + if stripped.startswith("{") or stripped.startswith("["): + return parse_policy_json(text) + return parse_policy_config(text) def audit_policies(policies: list[PolicyRule]) -> list[Finding]: diff --git a/tests/test_policies.py b/tests/test_policies.py index 3bacfde..1469378 100644 --- a/tests/test_policies.py +++ b/tests/test_policies.py @@ -1,6 +1,6 @@ import unittest -from fgai.policies import audit_policies, parse_policy_config +from fgai.policies import audit_policies, parse_policy_config, parse_policy_json class PolicyTests(unittest.TestCase): @@ -25,6 +25,29 @@ end self.assertIn("Broad allow policy", titles) self.assertIn("Accepted traffic lacks UTM inspection", titles) + def test_policy_audit_reads_fortios_api_json_shape(self): + policies = parse_policy_json( + """ +{ + "results": [ + { + "policyid": 7, + "srcaddr": [{"name": "all"}], + "dstaddr": [{"name": "all"}], + "service": [{"name": "ALL"}], + "action": "accept", + "logtraffic": "disable" + } + ] +} +""" + ) + + findings = audit_policies(policies) + + self.assertEqual(policies[0].policy_id, "7") + self.assertIn("Broad allow policy", {finding.title for finding in findings}) + if __name__ == "__main__": unittest.main()