update

src/fgai/cli.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import argparse
 import json
 import sys
+from pathlib import Path
 
 from .llm import ollama_summary
 from .logs import read_events, summarize_events
@@ -70,6 +71,28 @@ def suggest_blocks(args: argparse.Namespace) -> int:
     return 0
 
 
+def test_connection(args: argparse.Namespace) -> int:
+    client = FortiGateClient.from_env()
+    status = client.system_status()
+    results = status.get("results", status)
+    print("FortiGate API connection OK")
+    _print_json(results)
+    return 0
+
+
+def fetch_policies(args: argparse.Namespace) -> int:
+    client = FortiGateClient.from_env()
+    policies = client.firewall_policies()
+    if args.output:
+        output = Path(args.output)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(policies, indent=2, sort_keys=True), encoding="utf-8")
+        print(f"Wrote {output}")
+    else:
+        _print_json(policies)
+    return 0
+
+
 def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Local FortiGate AI/ML inspection tool")
     subparsers = parser.add_subparsers(required=True)
@@ -98,6 +121,13 @@ def build_parser() -> argparse.ArgumentParser:
     blocks.add_argument("--expiry-minutes", type=int, default=60)
     blocks.set_defaults(func=suggest_blocks)
 
+    connection = subparsers.add_parser("test-connection", help="Test FortiGate REST API credentials")
+    connection.set_defaults(func=test_connection)
+
+    fetch = subparsers.add_parser("fetch-policies", help="Fetch firewall policies through the FortiGate REST API")
+    fetch.add_argument("--output", help="Write JSON response to this file")
+    fetch.set_defaults(func=fetch_policies)
+
     return parser
 
 

src/fgai/mitigation.py

@@ -8,6 +8,7 @@ from typing import Iterable
 from urllib import error, request
 import json
 import ssl
+from urllib.parse import urlencode
 
 from .logs import event_score, is_utm_event
 from .models import BlockCandidate, LogEvent
@@ -107,3 +108,32 @@ class FortiGateClient:
             raise RuntimeError(f"FortiGate API returned HTTP {exc.code}: {detail}") from exc
         except error.URLError as exc:
             raise RuntimeError(f"Could not reach FortiGate API: {exc.reason}") from exc
+
+    def get_json(self, path: str, params: dict[str, str] | None = None) -> dict[str, object]:
+        query = f"?{urlencode(params)}" if params else ""
+        url = f"https://{self.host}{path}{query}"
+        req = request.Request(
+            url,
+            method="GET",
+            headers={
+                "Authorization": f"Bearer {self.api_token}",
+                "Accept": "application/json",
+            },
+        )
+        context = None if self.verify_tls else ssl._create_unverified_context()
+        try:
+            with request.urlopen(req, timeout=15, context=context) as response:
+                return json.loads(response.read().decode("utf-8"))
+        except error.HTTPError as exc:
+            detail = exc.read().decode("utf-8", errors="replace")
+            raise RuntimeError(f"FortiGate API returned HTTP {exc.code}: {detail}") from exc
+        except error.URLError as exc:
+            raise RuntimeError(f"Could not reach FortiGate API: {exc.reason}") from exc
+        except json.JSONDecodeError as exc:
+            raise RuntimeError("FortiGate API returned invalid JSON") from exc
+
+    def system_status(self) -> dict[str, object]:
+        return self.get_json("/api/v2/monitor/system/status")
+
+    def firewall_policies(self) -> dict[str, object]:
+        return self.get_json("/api/v2/cmdb/firewall/policy")