secagentx 3.0.0__py3-none-any.whl

backend/tools/base.py

@@ -0,0 +1,445 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Optional
+import os
+
+
+@dataclass
+class ToolResult:
+    success: bool = True
+    data: Any = None
+    error: str = ""
+    duration_ms: float = 0.0
+
+
+class BaseTool(ABC):
+    name: str = ""
+    description: str = ""
+    parameters: dict = None  # 子类在类级别重写此字段
+    agent_id: str = ""
+
+    @property
+    def params(self) -> dict:
+        """安全访问 parameters，兼容子类类级别定义和实例级别定义"""
+        val = self.parameters
+        if val is None:
+            return {"type": "object", "properties": {}}
+        return val
+
+    @abstractmethod
+    async def execute(self, **kwargs) -> ToolResult:
+        ...
+
+    def to_openai_function(self) -> dict:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.params,
+            },
+        }
+
+    def to_tool_call(self) -> dict:
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": self.params,
+        }
+
+
+# ═══════════════════════════════════════════════════════════════
+# 防火墙适配器接口 — 支持多种后端实现（企业级完善）
+# ═══════════════════════════════════════════════════════════════
+
+class FirewallAdapter(ABC):
+    """防火墙后端适配器抽象接口"""
+
+    @abstractmethod
+    async def block_ip(self, ip: str, reason: str = "",
+                        duration_minutes: int = 120) -> tuple[bool, str]:
+        ...
+
+    @abstractmethod
+    async def unblock_ip(self, ip: str) -> tuple[bool, str]:
+        ...
+
+    @abstractmethod
+    async def list_rules(self) -> list[dict]:
+        ...
+
+    @abstractmethod
+    async def check_ip(self, ip: str) -> dict:
+        ...
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        ...
+
+
+class IptablesAdapter(FirewallAdapter):
+    """iptables/nftables 真实防火墙后端（生产可用）
+
+    生产环境注意事项（已实测验证）:
+      - 不使用 `-m comment` 模块：在 nf_tables 后端下，添加时带 comment
+        删除时匹配不一致，会导致规则残留（实测 Bug）
+      - 使用独立链 `SecAgentX-BLOCK`：封禁规则全部放到该链，
+        list_rules 只查该链，不会与 INPUT 链上的其他规则混淆
+      - 封禁/解封需要 root 权限：通过 sudo 执行
+      - 删除规则时带完整匹配条件（-s ip -j DROP），避免误删其他规则
+    """
+
+    CHAIN = "SecAgentX-BLOCK"
+
+    def __init__(self, chain: str = None, use_sudo: bool = True):
+        self.chain = chain or self.CHAIN
+        # 生产环境 iptables 通常需要 root；可配置为 False（已是 root）
+        self.use_sudo = use_sudo or os.geteuid() != 0
+
+    async def _run_iptables(self, args: list[str]) -> tuple[int, str, str]:
+        import asyncio
+        import subprocess
+        cmd = ["sudo", "iptables", *args] if self.use_sudo else ["iptables", *args]
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+        )
+        stdout, stderr = await proc.communicate()
+        return proc.returncode or 0, stdout.decode(), stderr.decode()
+
+    async def _ensure_chain(self):
+        """确保独立链存在（不存在则创建并接入 INPUT）。"""
+        rc, stdout, _ = await self._run_iptables(["-L", self.chain, "-n"])
+        if rc != 0:
+            # 创建独立链
+            await self._run_iptables(["-N", self.chain])
+            # 接入 INPUT 链（在 ACCEPT 之前插入）
+            await self._run_iptables(["-I", "INPUT", "-j", self.chain])
+
+    async def block_ip(self, ip: str, reason: str = "",
+                        duration_minutes: int = 120) -> tuple[bool, str]:
+        await self._ensure_chain()
+        rc, _, stderr = await self._run_iptables([
+            "-A", self.chain, "-s", ip, "-j", "DROP",
+        ])
+        if rc != 0:
+            return False, f"iptables block failed: {stderr}"
+        return True, ""
+
+    async def unblock_ip(self, ip: str) -> tuple[bool, str]:
+        # 在独立链中删除，带完整匹配
+        rc, _, stderr = await self._run_iptables([
+            "-D", self.chain, "-s", ip, "-j", "DROP",
+        ])
+        if rc != 0:
+            # 规则可能不存在，不是致命错误
+            return True, f"rule not found (already removed): {stderr}"
+        return True, ""
+
+    async def list_rules(self) -> list[dict]:
+        rc, stdout, stderr = await self._run_iptables([
+            "-L", self.chain, "-n",
+        ])
+        if rc != 0:
+            return []
+        rules = []
+        for line in stdout.splitlines():
+            parts = line.split()
+            # 格式: DROP all -- <src> <dst>
+            if len(parts) >= 4 and parts[0] == "DROP":
+                ip = parts[3]
+                if ip != "0.0.0.0/0":
+                    rules.append({"ip": ip, "reason": ""})
+        return rules
+
+    async def check_ip(self, ip: str) -> dict:
+        rules = await self.list_rules()
+        for r in rules:
+            if r["ip"] == ip:
+                return {"is_blocked": True, "rule": r}
+        return {"is_blocked": False}
+
+    async def health_check(self) -> bool:
+        # 确保链存在（不存在则创建并接入 INPUT）——成功即健康
+        try:
+            await self._ensure_chain()
+            rc, _, stderr = await self._run_iptables(["-L", self.chain, "-n"])
+            return rc == 0 and "Permission" not in stderr
+        except Exception:
+            return False
+
+
+class DisabledFirewallAdapter(FirewallAdapter):
+    """安全默认值：不执行也不伪造任何防火墙变更。"""
+
+    _ERROR = "防火墙后端未配置；请显式设置 FIREWALL_BACKEND"
+
+    async def block_ip(self, ip: str, reason: str = "",
+                       duration_minutes: int = 120) -> tuple[bool, str]:
+        return False, self._ERROR
+
+    async def unblock_ip(self, ip: str) -> tuple[bool, str]:
+        return False, self._ERROR
+
+    async def list_rules(self) -> list[dict]:
+        return []
+
+    async def check_ip(self, ip: str) -> dict:
+        return {"is_blocked": False, "status": "disabled"}
+
+    async def health_check(self) -> bool:
+        return False
+
+
+class MockFirewallAdapter(FirewallAdapter):
+    """本地 JSON 测试后端；生产不得启用。"""
+
+    # 文件写操作锁 — 防止并发写入导致数据竞态
+    _file_lock = None
+
+    def __init__(self):
+        self._rules: dict[str, dict] = {}
+        if MockFirewallAdapter._file_lock is None:
+            import asyncio
+            MockFirewallAdapter._file_lock = asyncio.Lock()
+        self._load()
+
+    def _load(self):
+        import json
+        path = os.getenv("BLACKLIST_FILE", "data/blacklist/blacklist.json")
+        if os.path.exists(path):
+            try:
+                with open(path) as f:
+                    self._rules = json.load(f)
+            except Exception:
+                self._rules = {}
+
+    async def _save(self):
+        import json
+        path = os.getenv("BLACKLIST_FILE", "data/blacklist/blacklist.json")
+        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+        tmp = path + ".tmp"
+        async with MockFirewallAdapter._file_lock:
+            with open(tmp, "w") as f:
+                json.dump(self._rules, f, indent=2)
+            os.replace(tmp, path)
+
+    async def block_ip(self, ip: str, reason: str = "",
+                        duration_minutes: int = 120) -> tuple[bool, str]:
+        from datetime import datetime, timezone, timedelta
+        self._rules[ip] = {
+            "reason": reason or "auto-block",
+            "blocked_at": datetime.now(timezone.utc).isoformat(),
+            "expire_at": (datetime.now(timezone.utc) + timedelta(minutes=duration_minutes)).isoformat(),
+            "duration_minutes": duration_minutes,
+        }
+        await self._save()
+        return True, ""
+
+    async def unblock_ip(self, ip: str) -> tuple[bool, str]:
+        self._rules.pop(ip, None)
+        await self._save()
+        return True, ""
+
+    async def list_rules(self) -> list[dict]:
+        from datetime import datetime, timezone
+        now = datetime.now(timezone.utc)
+        active = []
+        for ip, rule in self._rules.items():
+            expire = datetime.fromisoformat(rule["expire_at"])
+            if expire > now:
+                active.append({"ip": ip, **rule})
+        return active
+
+    async def check_ip(self, ip: str) -> dict:
+        rules = await self.list_rules()
+        for r in rules:
+            if r["ip"] == ip:
+                return {"is_blocked": True, "rule": r}
+        return {"is_blocked": False}
+
+    async def health_check(self) -> bool:
+        return True
+
+
+class NftablesAdapter(FirewallAdapter):
+    """nftables 防火墙后端 — 现代 Linux 防火墙（iptables 的替代品）"""
+
+    def __init__(self, table: str = "secagentx", chain: str = "input"):
+        self.table = table
+        self.chain = chain
+        self._initialized = False
+        # 将初始化延迟到首次使用时（避免在 __init__ 中调用异步代码）
+
+    async def _ensure_table(self):
+        """确保 nftables 表和链存在（延迟初始化）"""
+        if self._initialized:
+            return
+        try:
+            import asyncio
+            import subprocess as _subprocess
+            proc = await asyncio.create_subprocess_exec(
+                "nft", "list", "table", "inet", self.table,
+                stdout=_subprocess.PIPE, stderr=_subprocess.PIPE,
+            )
+            await proc.communicate()
+            if proc.returncode != 0:
+                # 表不存在，创建
+                await asyncio.create_subprocess_exec(
+                    "nft", "add", "table", "inet", self.table,
+                )
+                await asyncio.create_subprocess_exec(
+                    "nft", "add", "chain", "inet", self.table, self.chain,
+                    "{ type filter hook input priority 0; policy accept; }",
+                )
+            self._initialized = True
+        except Exception:
+            pass  # 非 root 用户下跳过，实际封禁时会报错
+
+    async def _run_nft(self, args: list[str]) -> tuple[int, str, str]:
+        import asyncio, subprocess
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                "nft", *args,
+                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            )
+            stdout, stderr = await proc.communicate()
+            return proc.returncode or 0, stdout.decode(), stderr.decode()
+        except FileNotFoundError:
+            return 127, "", "nft executable is not installed"
+        except OSError as exc:
+            return 1, "", str(exc)
+
+    async def block_ip(self, ip: str, reason: str = "",
+                        duration_minutes: int = 120) -> tuple[bool, str]:
+        await self._ensure_table()
+        comment = f"SecAgentX: {reason[:60]}" if reason else "SecAgentX: auto-block"
+        rc, _, stderr = await self._run_nft([
+            "add", "rule", "inet", self.table, self.chain,
+            "ip", "saddr", ip, "drop",
+            "comment", f'"{comment}"',
+        ])
+        if rc != 0:
+            return False, f"nftables block failed: {stderr}"
+        return True, ""
+
+    async def unblock_ip(self, ip: str) -> tuple[bool, str]:
+        await self._ensure_table()
+        # nftables 没有 "delete rule by match" 的直接方式，需要 handle
+        # 先列出所有规则，找到对应 handle，再删除
+        rc, stdout, stderr = await self._run_nft([
+            "-a", "list", "chain", "inet", self.table, self.chain,
+        ])
+        if rc != 0:
+            return True, f"cannot list rules (may not exist): {stderr}"
+
+        import re
+        for line in stdout.splitlines():
+            if f"saddr {ip}" in line and "drop" in line:
+                m = re.search(r'# handle (\d+)', line)
+                if m:
+                    handle = m.group(1)
+                    rc2, _, stderr2 = await self._run_nft([
+                        "delete", "rule", "inet", self.table, self.chain,
+                        "handle", handle,
+                    ])
+                    if rc2 == 0:
+                        return True, ""
+                    return False, f"nftables unblock failed: {stderr2}"
+        return True, "rule not found (already removed)"
+
+    async def list_rules(self) -> list[dict]:
+        await self._ensure_table()
+        import re
+        rc, stdout, stderr = await self._run_nft([
+            "-a", "list", "chain", "inet", self.table, self.chain,
+        ])
+        if rc != 0:
+            return []
+        rules = []
+        for line in stdout.splitlines():
+            if "saddr" in line and "drop" in line:
+                parts = line.split()
+                ip = ""
+                for i, p in enumerate(parts):
+                    if p == "saddr" and i + 1 < len(parts):
+                        ip = parts[i + 1]
+                comment_m = re.search(r'comment "([^"]*)"', line)
+                handle_m = re.search(r'# handle (\d+)', line)
+                rules.append({
+                    "ip": ip,
+                    "reason": comment_m.group(1).replace("SecAgentX: ", "") if comment_m else "",
+                    "handle": handle_m.group(1) if handle_m else "",
+                })
+        return rules
+
+    async def check_ip(self, ip: str) -> dict:
+        await self._ensure_table()
+        rules = await self.list_rules()
+        for r in rules:
+            if r["ip"] == ip:
+                return {"is_blocked": True, "rule": r}
+        return {"is_blocked": False}
+
+    async def health_check(self) -> bool:
+        await self._ensure_table()
+        rc, _, _ = await self._run_nft(["list", "tables"])
+        return rc == 0
+
+
+class FirewallAdapterFactory:
+    """防火墙适配器工厂 — 通过环境变量切换后端"""
+
+    BACKENDS = {
+        "disabled": DisabledFirewallAdapter,
+        "mock": MockFirewallAdapter,
+        "iptables": IptablesAdapter,
+        "nftables": NftablesAdapter,
+        "aliyun": None,    # 由 _lazy_load 动态加载
+        "tencent": None,   # 由 _lazy_load 动态加载
+        "aws": None,       # 由 _lazy_load 动态加载
+    }
+
+    _LAZY_LOADERS = {
+        "aliyun": {
+            "module": "backend.tools.aliyun_firewall",
+            "class": "AliyunFirewallAdapter",
+            "pip": "alibabacloud_ecs20140526",
+        },
+        "tencent": {
+            "module": "backend.tools.tencent_firewall",
+            "class": "TencentFirewallAdapter",
+            "pip": "tencentcloud-sdk-python-vpc",
+        },
+        "aws": {
+            "module": "backend.tools.aws_firewall",
+            "class": "AWSFirewallAdapter",
+            "pip": "boto3",
+        },
+    }
+
+    @classmethod
+    def create(cls, backend: str = None) -> FirewallAdapter:
+        backend = backend or os.getenv("FIREWALL_BACKEND", "disabled")
+        adapter_class = cls.BACKENDS.get(backend)
+
+        # 延迟加载云厂商适配器
+        if adapter_class is None and backend in cls._LAZY_LOADERS:
+            loader = cls._LAZY_LOADERS[backend]
+            try:
+                import importlib
+                mod = importlib.import_module(loader["module"])
+                adapter_class = getattr(mod, loader["class"])
+                cls.BACKENDS[backend] = adapter_class
+            except ImportError as e:
+                raise ValueError(
+                    f"{loader['class']} 加载失败: {e}\n"
+                    f"请安装: pip install {loader['pip']}"
+                )
+
+        if adapter_class is None:
+            raise ValueError(
+                f"不支持的防火墙后端: '{backend}'。"
+                f"可选: {', '.join(k for k, v in cls.BACKENDS.items() if v)}"
+            )
+        return adapter_class()

backend/tools/calling.py

@@ -0,0 +1,207 @@
+"""
+统一 Tool Calling 核心模块
+
+标准化的 ToolCall 请求/响应数据模型，
+以及 LLM Function Calling 格式 ↔ 内部工具系统的适配层。
+"""
+import json
+import uuid
+import time
+from dataclasses import dataclass, field
+from typing import Any, Optional
+from .base import BaseTool, ToolResult
+
+
+# ═══════════════════════ 标准化数据模型 ═══════════════════════
+
+@dataclass
+class ToolCall:
+    """
+    一次工具调用的完整请求记录
+
+    由 LLM 的 tool_calls 解析而来。
+    """
+    call_id: str = field(default="")       # 对应 LLM 返回的 tool_call_id
+    tool_name: str = field(default="")     # 工具名（如 "threat_intel"）
+    arguments: dict = field(default_factory=dict)  # 解析后的参数字典
+    raw_arguments: str = ""               # LLM 返回的原始 JSON 字符串
+    timestamp: float = 0.0                # 创建时间
+    source: str = ""                      # 来源标识（"planner" / "agent:analyst" / "react"）
+    round_number: int = 0                 # ReAct 轮次
+
+    def __post_init__(self):
+        if not self.timestamp:
+            self.timestamp = time.time()
+        if not self.call_id:
+            self.call_id = f"tc_{uuid.uuid4().hex[:12]}"
+
+
+@dataclass
+class ToolCallResult:
+    """一次工具调用的完整结果记录"""
+    call: ToolCall
+    success: bool = False
+    result: Any = None
+    error: str = ""
+    duration_ms: float = 0.0
+    formatted_for_llm: str = ""
+
+    def to_llm_message(self) -> dict:
+        """转换为 OpenAI 兼容的 tool 角色消息"""
+        content = self.formatted_for_llm or json.dumps(
+            {"success": self.success, "result": self.result, "error": self.error},
+            ensure_ascii=False,
+        )
+        return {
+            "role": "tool",
+            "tool_call_id": self.call.call_id,
+            "content": content,
+        }
+
+
+@dataclass
+class ToolCallHistory:
+    """完整工具调用历史（可用于审计和回放）"""
+    MAX_HISTORY = 5000  # 最多保留 5000 条，防止 OOM
+    calls: list[ToolCallResult] = field(default_factory=list)
+
+    def add(self, result: ToolCallResult):
+        self.calls.append(result)
+        if len(self.calls) > self.MAX_HISTORY:
+            # 保留最新的 80%
+            self.calls = self.calls[-self.MAX_HISTORY:]
+
+    def to_dict(self) -> list[dict]:
+        return [
+            {
+                "call_id": r.call.call_id,
+                "tool": r.call.tool_name,
+                "arguments": r.call.arguments,
+                "success": r.success,
+                "duration_ms": r.duration_ms,
+                "source": r.call.source,
+                "round": r.call.round_number,
+            }
+            for r in self.calls
+        ]
+
+    @property
+    def total_calls(self) -> int:
+        return len(self.calls)
+
+    @property
+    def failed_calls(self) -> list[ToolCallResult]:
+        return [r for r in self.calls if not r.success]
+
+
+# ═══════════════════════ LLM 适配函数 ═══════════════════════
+
+def build_tools_for_llm(tools: list[BaseTool]) -> list[dict]:
+    """
+    将内部工具列表转换为 OpenAI 兼容的 tools 参数格式
+
+    输出示例:
+    [{
+        "type": "function",
+        "function": {
+            "name": "threat_intel",
+            "description": "查询IP/域名/哈希的威胁情报",
+            "parameters": {"type": "object", "properties": {...}, "required": [...]}
+        }
+    }]
+    """
+    return [t.to_openai_function() for t in tools]
+
+
+def parse_tool_calls(
+    llm_response: dict,
+    source: str = "",
+    round_number: int = 0,
+) -> list[ToolCall]:
+    """
+    从 LLM 响应中解析 tool_calls
+
+    兼容:
+    - OpenAI 标准格式: choices[0].message.tool_calls[]
+    - DeepSeek 格式: 同上（兼容 OpenAI）
+    - Function_call 旧格式: choices[0].message.function_call（兼容遗留）
+
+    Args:
+        llm_response: LLM API 返回的完整 JSON 响应
+        source: 来源标识
+        round_number: ReAct 轮次
+
+    Returns:
+        list[ToolCall]: 解析后的调用列表（可能为空）
+    """
+    raw_calls = []
+
+    try:
+        message = llm_response["choices"][0]["message"]
+
+        # 标准 OpenAI 格式
+        if "tool_calls" in message and message["tool_calls"]:
+            raw_calls = message["tool_calls"]
+        # 兼容旧版 function_call（部分 provider 仍返回这个）
+        elif "function_call" in message and message["function_call"]:
+            fc = message["function_call"]
+            raw_calls = [{
+                "id": f"call_{uuid.uuid4().hex[:12]}",
+                "type": "function",
+                "function": {
+                    "name": fc.get("name", ""),
+                    "arguments": fc.get("arguments", "{}"),
+                },
+            }]
+    except (KeyError, IndexError, TypeError):
+        return []
+
+    tool_calls = []
+    for tc in raw_calls:
+        try:
+            raw_args = tc.get("function", {}).get("arguments", "{}")
+            if isinstance(raw_args, str):
+                arguments = json.loads(raw_args)
+            else:
+                arguments = raw_args
+
+            tool_calls.append(ToolCall(
+                call_id=tc.get("id", f"tc_{uuid.uuid4().hex[:12]}"),
+                tool_name=tc.get("function", {}).get("name", ""),
+                arguments=arguments,
+                raw_arguments=raw_args if isinstance(raw_args, str)
+                              else json.dumps(raw_args, ensure_ascii=False),
+                source=source,
+                round_number=round_number,
+            ))
+        except (json.JSONDecodeError, KeyError, TypeError):
+            # 参数解析失败——保留记录但 arguments 为空
+            tool_calls.append(ToolCall(
+                call_id=tc.get("id", f"tc_{uuid.uuid4().hex[:12]}"),
+                tool_name=tc.get("function", {}).get("name", ""),
+                arguments={},
+                raw_arguments=str(tc.get("function", {}).get("arguments", "{}")),
+                source=source,
+                round_number=round_number,
+            ))
+
+    return tool_calls
+
+
+def format_tool_result_for_llm(tool_name: str, result: ToolResult) -> str:
+    """
+    将工具执行结果格式化为 LLM 可读的文本
+
+    规则:
+    - 成功: 结构化展示关键字段，超过 3000 字符截断
+    - 失败: 清晰说明错误原因
+    """
+    if result.success:
+        data = result.data or {}
+        data_str = json.dumps(data, ensure_ascii=False, indent=2)
+        if len(data_str) > 3000:
+            data_str = data_str[:3000] + f"\n... (truncated, {len(data_str)} total)"
+        return f"工具 [{tool_name}] 执行成功 ({result.duration_ms:.0f}ms):\n{data_str}"
+    else:
+        return f"工具 [{tool_name}] 执行失败: {result.error}"
+