AI Agent 安全设计：权限模型、沙箱隔离与审计日志

from dataclasses import dataclass, field
from enum import Enum
from typing import Callable


class RiskLevel(Enum):
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4


@dataclass
class ToolPermission:
    tool_name: str
    allowed_scopes: list[str] = field(default_factory=list)
    max_risk_level: RiskLevel = RiskLevel.MEDIUM
    rate_limit_per_minute: int = 60
    requires_confirmation: bool = False


@dataclass
class AgentRole:
    name: str
    permissions: list[ToolPermission] = field(default_factory=list)
    max_concurrent_tools: int = 5
    data_access_level: int = 0


class PermissionEngine:
    def __init__(self):
        self.roles: dict[str, AgentRole] = {}
        self.abac_policies: list[Callable] = []

    def register_role(self, role: AgentRole):
        self.roles[role.name] = role

    def add_abac_policy(self, policy_fn: Callable):
        self.abac_policies.append(policy_fn)

    def check_permission(
        self,
        agent_role: str,
        tool_name: str,
        context: dict,
    ) -> tuple[bool, str]:
        role = self.roles.get(agent_role)
        if not role:
            return False, f"未知角色: {agent_role}"

        tool_perm = next(
            (p for p in role.permissions if p.tool_name == tool_name), None
        )
        if not tool_perm:
            return False, f"角色 {agent_role} 无权访问工具 {tool_name}"

        current_risk = context.get("risk_level", RiskLevel.LOW)
        if current_risk.value > tool_perm.max_risk_level.value:
            return False, f"当前风险等级 {current_risk.name} 超出工具允许上限"

        for policy_fn in self.abac_policies:
            allowed, reason = policy_fn(role, tool_perm, context)
            if not allowed:
                return False, reason

        return True, "权限检查通过"

class DynamicPermissionAdjuster:
    def __init__(self, perm_engine: PermissionEngine):
        self.engine = perm_engine
        self.anomaly_counters: dict[str, int] = {}

    def record_tool_call(self, agent_id: str, tool_name: str):
        key = f"{agent_id}:{tool_name}"
        self.anomaly_counters[key] = self.anomaly_counters.get(key, 0) + 1

    def adjust_risk_level(self, agent_id: str, context: dict) -> RiskLevel:
        base_risk = context.get("base_risk", RiskLevel.LOW)

        write_call_count = sum(
            count for key, count in self.anomaly_counters.items()
            if key.startswith(agent_id) and ":write_" in key
        )
        if write_call_count > 20:
            return RiskLevel.CRITICAL
        if write_call_count > 10:
            return RiskLevel.HIGH
        if write_call_count > 5:
            return RiskLevel.MEDIUM

        return base_risk

import subprocess
import json
from pathlib import Path


class CodeSandbox:
    def __init__(self, sandbox_type: str = "docker"):
        self.sandbox_type = sandbox_type

    def execute(
        self,
        code: str,
        language: str = "python",
        timeout_seconds: int = 30,
        memory_limit_mb: int = 256,
    ) -> dict:
        if self.sandbox_type == "docker":
            return self._execute_docker(code, language, timeout_seconds, memory_limit_mb)
        elif self.sandbox_type == "e2b":
            return self._execute_e2b(code, language, timeout_seconds)
        else:
            raise ValueError(f"不支持的沙箱类型: {self.sandbox_type}")

    def _execute_docker(
        self, code: str, language: str, timeout: int, mem_limit: int
    ) -> dict:
        image_map = {
            "python": "python:3.12-slim",
            "javascript": "node:20-slim",
        }
        image = image_map.get(language, image_map["python"])

        code_file = Path("/tmp/sandbox_input.py")
        code_file.write_text(code)

        cmd = [
            "docker", "run", "--rm",
            "--network=none",
            "--read-only",
            f"--memory={mem_limit}m",
            f"--cpus=0.5",
            "--pids-limit=64",
            "--cap-drop=ALL",
            "-v", f"{code_file}:/code/input.py:ro",
            image,
            "python", "/code/input.py",
        ]

        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=timeout
        )
        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "exit_code": result.returncode,
        }

    def _execute_e2b(self, code: str, language: str, timeout: int) -> dict:
        from e2b_code_interpreter import Sandbox

        sandbox = Sandbox()
        try:
            execution = sandbox.run_code(code)
            return {
                "stdout": execution.text or "",
                "stderr": execution.error or "",
                "exit_code": 0 if not execution.error else 1,
            }
        finally:
            sandbox.kill()

import json
import jsonschema
import re
from typing import Any


class ToolInputValidator:
    def __init__(self):
        self.schemas: dict[str, dict] = {}
        self.custom_validators: dict[str, callable] = {}

    def register_schema(self, tool_name: str, schema: dict):
        self.schemas[tool_name] = schema

    def register_validator(self, tool_name: str, validator_fn: callable):
        self.custom_validators[tool_name] = validator_fn

    def validate(self, tool_name: str, arguments: dict) -> tuple[bool, str]:
        schema = self.schemas.get(tool_name)
        if schema:
            try:
                jsonschema.validate(arguments, schema)
            except jsonschema.ValidationError as e:
                return False, f"参数校验失败: {e.message}"

        custom_fn = self.custom_validators.get(tool_name)
        if custom_fn:
            ok, reason = custom_fn(arguments)
            if not ok:
                return False, reason

        return True, "验证通过"


def sql_validator(args: dict) -> tuple[bool, str]:
    if "query" in args:
        query = args["query"].upper()
        dangerous_keywords = ["DROP", "TRUNCATE", "DELETE", "ALTER", "INSERT", "UPDATE"]
        for kw in dangerous_keywords:
            if kw in query:
                return False, f"禁止执行包含 {kw} 的 SQL 语句"
    return True, ""


def filesystem_validator(args: dict) -> tuple[bool, str]:
    allowed_base = "/workspace/data"
    if "path" in args:
        target = str(Path(args["path"]).resolve())
        if not target.startswith(allowed_base):
            return False, f"路径 {target} 超出允许范围 {allowed_base}"
    return True, ""


validator = ToolInputValidator()
validator.register_validator("sql_query", sql_validator)
validator.register_validator("read_file", filesystem_validator)

import re


class OutputFilter:
    PATTERNS = {
        "credit_card": r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",
        "email": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
        "phone_cn": r"\b1[3-9]\d{9}\b",
        "id_card_cn": r"\b\d{17}[\dXx]\b",
        "ip_private": r"\b(10\.\d{1,3}|172\.(1[6-9]|2\d|3[01])|192\.168)\.\d{1,3}\.\d{1,3}\b",
    }

    def __init__(self, mask_pii: bool = True, allowed_patterns: list[str] | None = None):
        self.mask_pii = mask_pii
        self.allowed_patterns = allowed_patterns

    def filter_output(self, raw_output: str, tool_name: str) -> str:
        if not self.mask_pii:
            return raw_output

        filtered = raw_output
        for pii_type, pattern in self.PATTERNS.items():
            if self.allowed_patterns and pii_type in self.allowed_patterns:
                continue
            filtered = re.sub(pattern, f"[{pii_type.upper()}_MASKED]", filtered)

        return filtered

    def filter_for_model(self, raw_output: str, max_length: int = 8000) -> str:
        filtered = self.filter_output(raw_output)
        if len(filtered) > max_length:
            filtered = filtered[:max_length] + "\n...[输出截断]"
        return filtered

import time
from collections import defaultdict


class ToolRateLimiter:
    def __init__(self):
        self.call_records: dict[str, list[float]] = defaultdict(list)

    def check_rate_limit(
        self, tool_name: str, max_calls: int = 30, window_seconds: int = 60
    ) -> bool:
        now = time.time()
        self.call_records[tool_name] = [
            t for t in self.call_records[tool_name]
            if now - t < window_seconds
        ]
        if len(self.call_records[tool_name]) >= max_calls:
            return False
        self.call_records[tool_name].append(now)
        return True


class ToolScope:
    FILE_SYSTEM = "filesystem"
    NETWORK = "network"
    DATABASE = "database"
    EXTERNAL_API = "external_api"


SCOPE_RESTRICTIONS = {
    ToolScope.FILE_SYSTEM: {
        "allowed_paths": ["/workspace", "/tmp"],
        "blocked_operations": ["chmod", "chown", "symlink"],
    },
    ToolScope.NETWORK: {
        "allowed_domains": ["api.openai.com", "api.anthropic.com"],
        "blocked_ports": [22, 23, 3389],
        "max_payload_bytes": 1024 * 1024,
    },
    ToolScope.DATABASE: {
        "allowed_operations": ["SELECT"],
        "blocked_operations": ["DROP", "TRUNCATE", "ALTER"],
        "max_rows_returned": 1000,
    },
}

from dataclasses import dataclass
from enum import IntEnum


class DataClassification(IntEnum):
    PUBLIC = 0
    INTERNAL = 1
    CONFIDENTIAL = 2
    RESTRICTED = 3


@dataclass
class DataField:
    name: str
    classification: DataClassification
    pii_type: str | None = None
    masking_rules: dict | None = None


class DataAccessControl:
    def __init__(self, agent_clearance: DataClassification = DataClassification.INTERNAL):
        self.agent_clearance = agent_clearance
        self.field_registry: dict[str, DataField] = {}

    def register_field(self, dataset: str, field: DataField):
        key = f"{dataset}.{field.name}"
        self.field_registry[key] = field

    def can_access_field(self, dataset: str, field_name: str) -> bool:
        key = f"{dataset}.{field_name}"
        field = self.field_registry.get(key)
        if not field:
            return False
        return field.classification <= self.agent_clearance

    def mask_field(self, dataset: str, field_name: str, value: str) -> str:
        key = f"{dataset}.{field_name}"
        field = self.field_registry.get(key)
        if not field:
            return "[ACCESS_DENIED]"

        if field.classification > self.agent_clearance:
            return "[CLASSIFIED]"

        if field.pii_type:
            return self._apply_pii_mask(value, field.pii_type)

        return value

    def _apply_pii_mask(self, value: str, pii_type: str) -> str:
        mask_map = {
            "name": value[0] + "*" * (len(value) - 1) if value else "",
            "phone": value[:3] + "****" + value[-4:] if len(value) >= 7 else "****",
            "email": value[0] + "***@" + value.split("@")[-1] if "@" in value else "***",
            "id_card": value[:4] + "**********" + value[-4:] if len(value) >= 14 else "****",
        }
        return mask_map.get(pii_type, "***")

class SecureQueryEngine:
    def __init__(self, dac: DataAccessControl):
        self.dac = dac

    def execute_query(
        self, dataset: str, fields: list[str], raw_data: list[dict]
    ) -> list[dict]:
        allowed_fields = [
            f for f in fields if self.dac.can_access_field(dataset, f)
        ]
        blocked_count = len(fields) - len(allowed_fields)

        results = []
        for row in raw_data:
            masked_row = {}
            for field_name in allowed_fields:
                raw_value = str(row.get(field_name, ""))
                masked_row[field_name] = self.dac.mask_field(
                    dataset, field_name, raw_value
                )
            results.append(masked_row)

        if blocked_count > 0:
            results[0]["_meta"] = f"{blocked_count} 个字段因权限不足被过滤"

        return results

import json
import time
import uuid
from typing import Any
from pathlib import Path


@dataclass
class AuditEvent:
    event_type: str
    agent_id: str
    session_id: str
    timestamp: float
    details: dict[str, Any]
    risk_level: str = "LOW"
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))

    def to_dict(self) -> dict:
        return {
            "event_id": self.event_id,
            "event_type": self.event_type,
            "agent_id": self.agent_id,
            "session_id": self.session_id,
            "timestamp": self.timestamp,
            "risk_level": self.risk_level,
            "details": self.details,
        }


class AuditLogger:
    def __init__(self, log_path: str = "/var/log/agent_audit"):
        self.log_path = Path(log_path)
        self.log_path.mkdir(parents=True, exist_ok=True)

    def log_event(self, event: AuditEvent):
        log_line = json.dumps(event.to_dict(), ensure_ascii=False) + "\n"
        date_str = time.strftime("%Y-%m-%d")
        log_file = self.log_path / f"audit_{date_str}.jsonl"

        with open(log_file, "a", encoding="utf-8") as f:
            f.write(log_line)

    def log_tool_call(
        self, agent_id: str, session_id: str,
        tool_name: str, arguments: dict, risk_level: str = "LOW"
    ):
        self.log_event(AuditEvent(
            event_type="tool_call",
            agent_id=agent_id,
            session_id=session_id,
            timestamp=time.time(),
            risk_level=risk_level,
            details={
                "tool_name": tool_name,
                "arguments_keys": list(arguments.keys()),
                "arguments_size": len(json.dumps(arguments)),
            },
        ))

    def log_permission_check(
        self, agent_id: str, session_id: str,
        tool_name: str, allowed: bool, reason: str
    ):
        self.log_event(AuditEvent(
            event_type="permission_check",
            agent_id=agent_id,
            session_id=session_id,
            timestamp=time.time(),
            risk_level="MEDIUM" if not allowed else "LOW",
            details={
                "tool_name": tool_name,
                "allowed": allowed,
                "reason": reason,
            },
        ))

@dataclass
class RiskScore:
    total: float
    components: dict[str, float]
    level: RiskLevel

    def to_dict(self) -> dict:
        return {
            "total": self.total,
            "level": self.level.name,
            "components": self.components,
        }


class RiskScoringEngine:
    WEIGHTS = {
        "tool_sensitivity": 0.30,
        "data_sensitivity": 0.25,
        "user_trust": 0.20,
        "behavior_anomaly": 0.15,
        "environment": 0.10,
    }

    TOOL_SENSITIVITY = {
        "read_file": 0.2,
        "write_file": 0.5,
        "sql_query": 0.6,
        "execute_code": 0.8,
        "delete_file": 0.9,
        "send_email": 0.7,
        "api_call_external": 0.6,
    }

    def compute_risk(self, context: dict) -> RiskScore:
        scores = {}

        tool = context.get("tool_name", "unknown")
        scores["tool_sensitivity"] = self.TOOL_SENSITIVITY.get(tool, 0.5)

        data_level = context.get("data_classification", 0)
        scores["data_sensitivity"] = data_level / 3.0

        trust = context.get("user_trust_score", 0.5)
        scores["user_trust"] = 1.0 - trust

        anomaly = context.get("anomaly_score", 0.0)
        scores["behavior_anomaly"] = min(anomaly, 1.0)

        env = context.get("is_production", False)
        scores["environment"] = 1.0 if env else 0.3

        total = sum(
            scores[k] * self.WEIGHTS[k] for k in self.WEIGHTS
        )

        if total >= 0.8:
            level = RiskLevel.CRITICAL
        elif total >= 0.6:
            level = RiskLevel.HIGH
        elif total >= 0.3:
            level = RiskLevel.MEDIUM
        else:
            level = RiskLevel.LOW

        return RiskScore(total=total, components=scores, level=level)

class HumanInTheLoopTrigger:
    def __init__(self, risk_engine: RiskScoringEngine):
        self.risk_engine = risk_engine
        self.auto_approve_max_level = RiskLevel.LOW
        self.require_confirm_levels = [RiskLevel.MEDIUM, RiskLevel.HIGH]
        self.block_levels = [RiskLevel.CRITICAL]

    def should_require_confirmation(
        self, context: dict
    ) -> tuple[bool, str, RiskScore]:
        risk = self.risk_engine.compute_risk(context)

        if risk.level in self.block_levels:
            return True, "操作已被安全策略阻止", risk

        if risk.level in self.require_confirm_levels:
            reason = (
                f"操作涉及 {context.get('tool_name', 'unknown')}，"
                f"风险等级: {risk.level.name}（总分: {risk.total:.2f}）"
            )
            return True, reason, risk

        return False, "", risk


class SecurityPolicyEngine:
    def __init__(self):
        self.risk_engine = RiskScoringEngine()
        self.human_loop = HumanInTheLoopTrigger(self.risk_engine)
        self.perm_engine = PermissionEngine()
        self.audit_logger = AuditLogger()
        self.rate_limiter = ToolRateLimiter()
        self.output_filter = OutputFilter(mask_pii=True)

    def evaluate_tool_call(
        self, agent_id: str, session_id: str,
        tool_name: str, arguments: dict,
        context: dict,
    ) -> dict:
        full_context = {**context, "tool_name": tool_name}

        if not self.rate_limiter.check_rate_limit(tool_name):
            self.audit_logger.log_event(AuditEvent(
                event_type="rate_limit_exceeded",
                agent_id=agent_id, session_id=session_id,
                timestamp=time.time(), risk_level="MEDIUM",
                details={"tool_name": tool_name},
            ))
            return {"allowed": False, "reason": "超出调用频率限制"}

        need_confirm, confirm_reason, risk = (
            self.human_loop.should_require_confirmation(full_context)
        )

        self.audit_logger.log_tool_call(
            agent_id, session_id, tool_name, arguments,
            risk_level=risk.level.name,
        )

        if need_confirm:
            return {
                "allowed": False,
                "need_confirmation": True,
                "reason": confirm_reason,
                "risk_score": risk.to_dict(),
            }

        return {"allowed": True, "risk_score": risk.to_dict()}

┌─────────────────────────────────────────────────────────────────────┐
│                    AI Agent 安全架构（分层视图）                       │
│                                                                     │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                     用户交互层                                 │  │
│  │   用户输入 → 输入清洗 → Prompt 注入检测 → 风险预评估             │  │
│  └───────────────────────────┬───────────────────────────────────┘  │
│                              ▼                                      │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                   安全策略引擎                                 │  │
│  │   ┌──────────────┐  ┌──────────────┐  ┌──────────────────┐   │  │
│  │   │  风险评分引擎  │  │  权限检查器   │  │  人工确认网关      │   │  │
│  │   │  RiskScoring  │  │  Permission   │  │  HITL Trigger    │   │  │
│  │   └──────────────┘  └──────────────┘  └──────────────────┘   │  │
│  └───────────────────────────┬───────────────────────────────────┘  │
│                              ▼                                      │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                   Agent 推理层                                 │  │
│  │   LLM 推理 → 工具选择决策 → 参数生成 → 安全约束注入              │  │
│  └───────────────────────────┬───────────────────────────────────┘  │
│                              ▼                                      │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                 工具调用沙箱层                                  │  │
│  │   输入验证 → 范围限制 → 速率控制 → 输出过滤                      │  │
│  │                                                               │  │
│  │   ┌──────────────────────────────────────────────────────┐    │  │
│  │   │                 代码执行沙箱                          │    │  │
│  │   │   Docker / gVisor / Firecracker / E2B               │    │  │
│  │   │   网络隔离 + 文件系统只读 + 资源限制 + 超时终止        │    │  │
│  │   └──────────────────────────────────────────────────────┘    │  │
│  └───────────────────────────┬───────────────────────────────────┘  │
│                              ▼                                      │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                 数据访问沙箱层                                  │  │
│  │   数据分级 → 字段级访问控制 → PII 检测 → 动态脱敏               │  │
│  └───────────────────────────┬───────────────────────────────────┘  │
│                              ▼                                      │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                   审计日志层                                   │  │
│  │   结构化日志 → 实时告警 → 合规报告 → 不可篡改存储               │  │
│  └───────────────────────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────────────────────┘