# Install Ollama
curl -fsSL https://ollama.com/install.sh | sh

# Pull Qwen2.5-7B (runs on 4GB VPS, Q4 quantized)
ollama pull qwen2.5:7b-instruct

# Verify
ollama run qwen2.5:7b-instruct "nftables flush ruleset syntax"
# Confirm the model outputs correct nft syntax snippets

#!/usr/bin/env python3
"""collect_rules.py — Export nftables ruleset + traffic profile"""

import subprocess
import json
from datetime import datetime


def get_ruleset() -> dict:
    """Export current nftables ruleset as JSON"""
    result = subprocess.run(
        ["nft", "-j", "list", "ruleset"],
        capture_output=True, text=True, check=True
    )
    return json.loads(result.stdout)


def get_conntrack_stats() -> str:
    """Get conntrack connection stats"""
    result = subprocess.run(
        ["conntrack", "-L", "-o", "export"],
        capture_output=True, text=True
    )
    return result.stdout


def get_sysctl_context() -> dict:
    """Get key kernel network parameters"""
    params = [
        "net.ipv4.ip_forward",
        "net.ipv4.conf.all.rpf_filter",
        "net.ipv4.conf.all.accept_redirects",
        "net.ipv4.conf.all.accept_source_route",
        "net.ipv4.conf.all.secure_redirects",
        "net.ipv4.tcp_syncookies",
    ]
    ctx = {}
    for p in params:
        try:
            r = subprocess.run(["sysctl", "-n", p], capture_output=True, text=True)
            ctx[p] = r.stdout.strip()
        except Exception:
            ctx[p] = "unknown"
    return ctx


def collect_all() -> dict:
    snapshot = {
        "timestamp": datetime.utcnow().isoformat(),
        "ruleset": get_ruleset(),
        "conntrack_sample": get_conntrack_stats(),
        "sysctl_context": get_sysctl_context(),
    }
    return snapshot


if __name__ == "__main__":
    data = collect_all()
    json.dump(data, open("/tmp/firewall-context.json", "w"), indent=2)
    print(f"Snapshot saved to /tmp/firewall-context.json")
    print(f"Ruleset size: {len(json.dumps(data['ruleset']))} bytes")

def collect_security_events(hours: int = 24) -> list:
    """Collect recent security events from fail2ban / audit logs"""
    events = []

    # fail2ban ban records
    try:
        import re
        with open("/var/log/fail2ban.log", "r") as f:
            for line in f:
                m = re.search(r"(?:Ban|Found) IP (\S+)", line)
                if m:
                    events.append({
                        "type": "fail2ban_ban",
                        "ip": m.group(1),
                        "timestamp": line.split()[0] if line else None,
                    })
    except FileNotFoundError:
        pass

    # Brute-force SSH attempts
    try:
        result = subprocess.run(
            ["journalctl", "-u", "sshd", "--since", f"{hours} hours ago",
             "--no-pager"],
            capture_output=True, text=True
        )
        for line in result.stdout.splitlines():
            if "Failed password" in line or "Invalid user" in line:
                events.append({
                    "type": "ssh_brute_force",
                    "raw": line[:200],
                })
    except Exception:
        pass

    return events[-50:]  # max 50 events

FIREWALL_SYSTEM_PROMPT = """\
You are a senior Linux network engineer specializing in nftables rule design.
Based on the provided context, generate a new nftables ruleset.

## Output format (strict JSON)
{
  "ruleset": { /* complete nftables JSON ruleset */ },
  "changes": [
    {"action": "add|remove|modify", "description": "...", "reason": "..."}
  ],
  "risk_assessment": {
    "ssh_accessible": true,
    "http_accessible": true,
    "breaks_existing_connections": false,
    "confidence": 0.95
  }
}

## Constraints
1. Must retain SSH (22/tcp) access for trusted CIDRs (see sysctl_context.trusted_cidrs)
2. Must retain 80/443 access for 0.0.0.0/0
3. Default DROP policy requires ESTABLISHED,RELATED accept rule
4. Must not delete lo interface rules
5. Every new rule must include a comment (nft # comment)
"""

def build_user_prompt(context: dict, security_events: list) -> str:
    parts = [
        "## Current nftables ruleset",
        json.dumps(context["ruleset"], indent=2),
        "",
        "## Kernel network parameters",
        json.dumps(context["sysctl_context"], indent=2),
        "",
        "## Recent connection stats (top sample)",
        context.get("conntrack_sample", "")[:2000],
        "",
    ]
    if security_events:
        parts.append("## Security events in last 24 hours")
        for ev in security_events[:20]:
            parts.append(f"- {ev['type']}: {ev.get('ip', ev.get('raw', ''))}")
        parts.append("")

    parts.append("Based on the above, generate an optimized nftables ruleset.")
    return "\n".join(parts)

import requests
import json
import time


def generate_rules(
    context: dict,
    security_events: list,
    model: str = "qwen2.5:7b-instruct",
    timeout: int = 120,
) -> dict:
    """Call local Ollama to generate nftables ruleset"""

    user_prompt = build_user_prompt(context, security_events)

    # Ollama is OpenAI-compatible
    resp = requests.post(
        "http://localhost:11434/v1/chat/completions",
        json={
            "model": model,
            "messages": [
                {"role": "system", "content": FIREWALL_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.1,   # low temp: rule generation needs determinism
            "max_tokens": 4096,
        },
        timeout=timeout,
    )
    resp.raise_for_status()

    content = resp.json()["choices"][0]["message"]["content"]

    # Extract JSON (LLM may wrap in markdown fences)
    content = content.strip()
    if content.startswith("```"):
        lines = content.split("\n")
        content = "\n".join(lines[1:-1]) if lines[-1].strip() == "```" else "\n".join(lines[1:])

    result = json.loads(content)
    result["generated_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ")
    return result

def ruleset_to_nft(result: dict) -> str:
    """Convert LLM output JSON to an nft command file"""
    # nft -j and nft -f both accept JSON format
    ruleset_json = json.dumps(result["ruleset"], indent=2)
    return ruleset_json

# Check syntax only, no kernel state change
nft -c -f new-ruleset.nft
echo "Exit code: $?"
# 0 = valid, non-zero = syntax error, abort

#!/bin/bash
# sandbox_test.sh — Test rules in an isolated network namespace
set -e

SANDBOX_NS="nft-sandbox-$$"

ip netns add "$SANDBOX_NS"
ip netns exec "$SANDBOX_NS" nft -f /tmp/new-ruleset.nft
sleep 2

ip netns exec "$SANDBOX_NS" bash -c '
  ping -c 1 127.0.0.1 > /dev/null || { echo "FAIL: lo unreachable"; exit 1; }
  echo "PASS: lo interface OK"
'

ip netns del "$SANDBOX_NS"
echo "Sandbox test passed"

def connectivity_smoke_test() -> bool:
    """Verify critical services reachable before applying rules"""
    checks = [
        ("localhost", 80, "HTTP service"),
        ("localhost", 443, "HTTPS service"),
        ("localhost", 22, "SSH service"),
    ]
    all_pass = True
    for host, port, desc in checks:
        try:
            import socket
            s = socket.create_connection((host, port), timeout=3)
            s.close()
            print(f"  ✓ {desc} ({host}:{port}) reachable")
        except Exception as e:
            print(f"  ✗ {desc} ({host}:{port}): {e}")
            all_pass = False
    return all_pass

#!/usr/bin/env python3
"""apply_rules.py — Atomic nftables reload with auto-rollback"""

import subprocess
import time
import os
from datetime import datetime


def atomic_apply(nft_file: str, timeout: int = 30) -> bool:
    """
    Atomic nftables reload:
    1. Snapshot current rules
    2. Load new rules
    3. Roll back on failure
    """
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = f"/tmp/nft-backup-{ts}.nft"

    # Step 1: Snapshot current rules
    print(f"[1/4] Snapshotting current ruleset → {backup_file}")
    subprocess.run(
        ["nft", "list", "ruleset"],
        stdout=open(backup_file, "w"),
        check=True,
    )
    print(f"      Backup size: {os.path.getsize(backup_file)} bytes")

    # Step 2: Load new rules
    print("[2/4] Loading new ruleset...")
    try:
        subprocess.run(["nft", "-f", nft_file], check=True, timeout=10)
        print("      New ruleset loaded successfully")
    except subprocess.CalledProcessError as e:
        print(f"      FAILED: {e}")
        return rollback(backup_file)

    # Step 3: Verify SSH still reachable
    print("[3/4] Verifying SSH still accessible...")
    time.sleep(1)
    try:
        import socket
        s = socket.create_connection(("127.0.0.1", 22), timeout=5)
        s.close()
        print("      SSH port 22 OK")
    except Exception as e:
        print(f"      SSH UNREACHABLE: {e}")
        print("      TRIGGERING AUTO-ROLLBACK")
        return rollback(backup_file)

    # Step 4: Verify HTTP
    print("[4/4] Verifying HTTP service...")
    try:
        import urllib.request
        urllib.request.urlopen("http://127.0.0.1:80", timeout=5)
        print("      HTTP OK")
    except Exception as e:
        print(f"      HTTP check: {e} (non-fatal)")

    print("✓ Ruleset applied successfully. Backup at:", backup_file)
    return True


def rollback(backup_file: str) -> bool:
    """Roll back to backup ruleset"""
    print(f"⚠  Rolling back to {backup_file}")
    try:
        subprocess.run(["nft", "-f", backup_file], check=True, timeout=15)
        print("✓ Rollback successful")
        return True
    except Exception as e:
        print(f"✗ ROLLBACK FAILED: {e}")
        print("  MANUAL INTERVENTION REQUIRED: nft -f", backup_file)
        return False


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: apply_rules.py <nft_file>")
        sys.exit(1)
    success = atomic_apply(sys.argv[1])
    sys.exit(0 if success else 1)

[Unit]
Description=LLM-driven nftables optimization
After=network-online.target ollama.service
Wants=network-online.target

[Service]
Type=oneshot
User=root
ExecStartPre=/usr/bin/python3 /opt/firewall-agent/collect_rules.py
ExecStartPre=/usr/bin/python3 /opt/firewall-agent/collect_events.py
ExecStart=/usr/bin/python3 /opt/firewall-agent/generate_and_apply.py
MemoryMax=1G
CPUQuota=50%
TimeoutStartSec=300

[Install]
WantedBy=multi-user.target

[Unit]
Description=Run LLM firewall optimization weekly

[Timer]
OnCalendar=Sun 03:00
RandomizedDelaySec=3600
Persistent=true

[Install]
WantedBy=timers.target

sudo systemctl enable --now llm-firewall-optimize.timer
sudo systemctl status llm-firewall-optimize.timer

SHOULD_OPTIMIZE_THRESHOLD = {
    "new_security_events": 5,    # more than 5 new banned IPs
    "rule_count_drift": 10,     # rule count drift > 10%
    "conntrack_overflow": True,  # conntrack table near limit
}

def should_optimize(context: dict, events: list) -> bool:
    new_bans = [e for e in events if e["type"] == "fail2ban_ban"]
    if len(new_bans) > SHOULD_OPTIMIZE_THRESHOLD["new_security_events"]:
        print(f"Trigger: {len(new_bans)} new banned IPs detected")
        return True
    return False

#!/usr/bin/env python3
"""
llm_firewall_agent.py — Complete LLM Firewall Optimization Agent
Runs on VPS, executes weekly via systemd timer.
"""

import json
import subprocess
import sys
from datetime import datetime
from pathlib import Path

LOG_DIR = Path("/var/log/llm-firewall")
WORK_DIR = Path("/tmp/llm-firewall")
OLLAMA_URL = "http://localhost:11434"
MODEL = "qwen2.5:7b-instruct"

LOG_DIR.mkdir(parents=True, exist_ok=True)
WORK_DIR.mkdir(parents=True, exist_ok=True)


def log(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line)
    (LOG_DIR / f"agent-{datetime.now().strftime('%Y%m%d')}.log").open("a").write(line + "\n")


def main():
    log("=== LLM Firewall Agent started ===")

    # Phase 1: Collect
    log("[Phase 1] Collecting current state...")
    context = json.loads(subprocess.check_output(
        ["python3", str(WORK_DIR / "collect_rules.py")]
    ).decode())

    # Phase 2: Generate
    log("[Phase 2] Generating new ruleset via LLM...")
    try:
        result = generate_rules(context, security_events=[])
        log(f"  Generated {len(result.get('changes', []))} changes")
    except Exception as e:
        log(f"  LLM generation failed: {e}")
        sys.exit(1)

    # Phase 3: Validate
    log("[Phase 3] Validating ruleset...")
    nft_file = WORK_DIR / "new-ruleset.nft"
    nft_file.write_text(ruleset_to_nft(result))

    check = subprocess.run(
        ["nft", "-c", "-f", str(nft_file)],
        capture_output=True, text=True
    )
    if check.returncode != 0:
        log(f"  Syntax check FAILED:\n{check.stderr}")
        sys.exit(1)
    log("  Syntax check PASSED")

    sandbox = subprocess.run(
        ["bash", str(WORK_DIR / "sandbox_test.sh")],
        capture_output=True, text=True
    )
    if sandbox.returncode != 0:
        log(f"  Sandbox test FAILED:\n{sandbox.stderr}")
        sys.exit(1)
    log("  Sandbox test PASSED")

    # Phase 4: Apply
    log("[Phase 4] Applying atomically...")
    apply_result = subprocess.run(
        ["python3", str(WORK_DIR / "apply_rules.py"), str(nft_file)],
        capture_output=True, text=True
    )
    if apply_result.returncode != 0:
        log(f"  Apply FAILED:\n{apply_result.stderr}")
        sys.exit(1)

    log("=== LLM Firewall Agent completed successfully ===")


if __name__ == "__main__":
    main()

# Check default policy is DROP
nft list ruleset | grep -A2 'policy drop'

# Check for 0.0.0.0/0 ACCEPT rules (dangerous)
nft list ruleset | grep -B2 'accept' | grep '0.0.0.0/0'