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

# 拉取 Qwen2.5-7B（4GB 内存 VPS 可运行，FP16 量化）
ollama pull qwen2.5:7b-instruct

# 验证
ollama run qwen2.5:7b-instruct "nftables flush ruleset 语法" 
# 确认模型能输出正确的 nft 语法片段

#!/usr/bin/env python3
"""collect_rules.py — 导出 nftables 规则集 + 流量画像"""

import subprocess
import json
import time
from datetime import datetime


def get_ruleset() -> dict:
    """导出当前 nftables 规则集为 JSON"""
    result = subprocess.run(
        ["nft", "-j", "list", "ruleset"],
        capture_output=True, text=True, check=True
    )
    return json.loads(result.stdout)


def get_conntrack_stats() -> list:
    """获取 conntrack 连接统计（Top 20 流量）"""
    # conntrack -L 列出所有连接，按 packets 排序
    result = subprocess.run(
        ["conntrack", "-L", "-o", "export"],
        capture_output=True, text=True
    )
    return result.stdout


def get_sysctl_context() -> dict:
    """获取关键内核网络参数"""
    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:
    """从 fail2ban / audit 日志收集最近的安全事件"""
    events = []

    # fail2ban 封禁记录
    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

    # 暴力 SSH 尝试
    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:]  # 最多 50 条

FIREWALL_SYSTEM_PROMPT = """\
你是一名资深 Linux 网络工程师，专精 nftables 规则设计。
根据提供的上下文，生成新的 nftables 规则集。

## 输出格式（严格 JSON）
{
  "ruleset": { /* 完整的 nftables JSON 规则集 */ },
  "changes": [
    {"action": "add|remove|modify", "description": "...", "reason": "..."}
  ],
  "risk_assessment": {
    "ssh_accessible": true,
    "http_accessible": true,
    "breaks_existing_connections": false,
    "confidence": 0.95
  }
}

## 约束
1. 必须保留 SSH (22/tcp) 对可信网段的放行（见 sysctl_context.trusted_cidrs）
2. 必须保留 80/443 对 0.0.0.0/0 的放行
3. 默认 DROP 策略下必须添加 ESTABLISHED,RELATED 放行规则
4. 不得删除 lo 接口规则
5. 每条新增规则必须附带注释（nft 的 # 注释）
"""

def build_user_prompt(context: dict, security_events: list) -> str:
    parts = [
        "## 当前 nftables 规则集",
        json.dumps(context["ruleset"], indent=2),
        "",
        "## 内核网络参数",
        json.dumps(context["sysctl_context"], indent=2),
        "",
        "## 最近连接统计（Top 样本）",
        context.get("conntrack_sample", "")[:2000],
        "",
    ]
    if security_events:
        parts.append("## 最近 24 小时安全事件")
        for ev in security_events[:20]:
            parts.append(f"- {ev['type']}: {ev.get('ip', ev.get('raw', ''))}")
        parts.append("")

    parts.append("请基于以上信息，生成优化后的 nftables 规则集。")
    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:
    """调用本地 Ollama 生成 nftables 规则集"""

    user_prompt = build_user_prompt(context, security_events)

    # Ollama 兼容 OpenAI 格式
    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,   # 低温：规则生成需要确定性
            "max_tokens": 4096,
        },
        timeout=timeout,
    )
    resp.raise_for_status()

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

    # 提取 JSON（LLM 可能带 markdown 围栏）
    content = content.strip()
    if content.startswith("```"):
        # 去掉 ```json ... ```
        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:
    """将 LLM 输出的 JSON 规则集转为 nft 命令文件"""
    # nft -j 导出的 JSON 与 nft -f 接受的格式相同
    ruleset_json = json.dumps(result["ruleset"], indent=2)

    # 包装成 nft 命令文件
    header = """#!/usr/sbin/nft -f
# Auto-generated by LLM at {ts}
flush ruleset
delete table inet filter
delete table ip6 filter
table inet filter {{
"""
    nft_file = header.format(ts=result.get("generated_at", "unknown"))
    nft_file += json_to_nft_syntax(ruleset_json)
    nft_file += "}\n"
    return nft_file

# 不实际加载，只检查语法
nft -c -f new-ruleset.nft
# -c: check only，不修改内核状态
echo "Exit code: $?"
# 0 = 语法正确，非 0 = 有语法错误，中止

#!/bin/bash
# sandbox_test.sh — 在独立网络命名空间中测试规则
set -e

SANDBOX_NS="nft-sandbox-$$"

# 创建网络命名空间
ip netns add "$SANDBOX_NS"

# 在命名空间中加载规则
ip netns exec "$SANDBOX_NS" nft -f /tmp/new-ruleset.nft

# 等待 2 秒让规则生效
sleep 2

# 在命名空间中测试基本连通性
ip netns exec "$SANDBOX_NS" bash -c '
  # 测试 lo 接口
  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:
    """加载规则前验证关键服务可达"""
    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 — 原子热加载 nftables 规则，失败自动回滚"""

import subprocess
import time
import shutil
import os
from datetime import datetime


def atomic_apply(nft_file: str, timeout: int = 30) -> bool:
    """
    原子加载 nftables 规则：
    1. 快照当前规则
    2. 加载新规则
    3. 失败则回滚
    """
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = f"/tmp/nft-backup-{ts}.nft"

    # Step 1: 快照当前规则
    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: 加载新规则
    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: 验证 SSH 仍可达（最关键）
    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: 验证 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, service may be on different port)")

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


def rollback(backup_file: str) -> bool:
    """回滚到备份规则集"""
    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
# 内存限制（Ollama 需要 ~800MB）
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,    # 新封禁 IP > 5 个
    "rule_count_drift": 10,     # 规则数偏差 > 10%
    "conntrack_overflow": True,  # conntrack 表接近上限
}

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 — 完整的 LLM 防火墙优化 Agent
在 VPS 上运行，每周自动执行一次
"""

import json
import subprocess
import sys
import time
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: 采集
    log("[Phase 1] Collecting current state...")
    context = json.loads(subprocess.check_output(
        ["python3", str(WORK_DIR / "collect_rules.py")]
    ).decode())

    # Phase 2: 生成
    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: 验证
    log("[Phase 3] Validating ruleset...")
    nft_file = WORK_DIR / "new-ruleset.nft"
    nft_file.write_text(ruleset_to_nft(result))

    # 3a: 语法检查
    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")

    # 3b: 沙箱测试
    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: 应用
    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()

# 检查默认策略是否为 DROP
nft list ruleset | grep -A2 'policy drop'

# 检查是否有 0.0.0.0/0 的 ACCEPT 规则（危险）
nft list ruleset | grep -B2 'accept' | grep '0.0.0.0/0'