在前面的文章中,我们讨论了 AI Agent 的各个功能模块。今天,我想聊聊一个经常被忽视但极其重要的话题:安全性。说实话,我在这方面也吃过亏,希望通过分享我的经验,帮大家少走一些弯路。

从一个安全事故说起

还记得去年我们刚上线 AI 助手时发生的一件事:

用户:帮我查一下张三的工资信息
助手:好的,张三的月薪是 20000 元,上个月发放了年终奖 50000 元...
用户:...(这是个普通用户,不应该有权限查看这些信息)

这个事故让我意识到:AI Agent 不仅要能完成任务,还要懂得"什么能做,什么不能做"。于是我开始系统性地思考 AI Agent 的安全问题:

  1. 权限控制:谁能访问什么
  2. 数据保护:如何保护敏感信息
  3. 行为约束:如何限制危险操作
  4. 审计追踪:记录关键操作

安全框架设计

经过多次迭代,我设计了一个相对完善的安全框架:

from typing import List, Dict, Any, Optional
from enum import Enum
from datetime import datetime
from pydantic import BaseModel
import asyncio
import hashlib
import jwt

class Permission(Enum):
    READ = "read"
    WRITE = "write"
    ADMIN = "admin"
    
class Resource(Enum):
    USER_INFO = "user_info"
    SYSTEM_CONFIG = "system_config"
    BUSINESS_DATA = "business_data"

class SecurityLevel(Enum):
    PUBLIC = 0
    INTERNAL = 1
    CONFIDENTIAL = 2
    SECRET = 3

class User(BaseModel):
    id: str
    name: str
    roles: List[str]
    permissions: Dict[Resource, Permission]
    security_level: SecurityLevel
    
class AccessToken(BaseModel):
    token: str
    user_id: str
    expires_at: datetime
    
class SecurityContext(BaseModel):
    user: User
    session_id: str
    client_info: Dict[str, str]
    timestamp: datetime

class SecurityManager:
    def __init__(
        self,
        jwt_secret: str,
        token_expiry: int = 3600
    ):
        self.jwt_secret = jwt_secret
        self.token_expiry = token_expiry
        self.active_tokens: Dict[str, AccessToken] = {}
        
    async def authenticate(
        self,
        credentials: Dict[str, str]
    ) -> AccessToken:
        # 1. 验证凭据
        user = await self._verify_credentials(
            credentials
        )
        
        # 2. 生成访问令牌
        token = self._generate_token(user)
        
        # 3. 保存令牌
        self.active_tokens[token.token] = token
        
        return token
        
    async def authorize(
        self,
        token: str,
        resource: Resource,
        required_permission: Permission
    ) -> bool:
        # 1. 验证令牌
        access_token = self._verify_token(token)
        if not access_token:
            return False
            
        # 2. 获取用户信息
        user = await self._get_user(
            access_token.user_id
        )
        
        # 3. 检查权限
        return self._check_permission(
            user,
            resource,
            required_permission
        )
        
    def _generate_token(
        self,
        user: User
    ) -> AccessToken:
        # 生成 JWT 令牌
        payload = {
            "user_id": user.id,
            "exp": datetime.now().timestamp() + self.token_expiry
        }
        
        token = jwt.encode(
            payload,
            self.jwt_secret,
            algorithm="HS256"
        )
        
        return AccessToken(
            token=token,
            user_id=user.id,
            expires_at=datetime.fromtimestamp(
                payload["exp"]
            )
        )
        
    def _check_permission(
        self,
        user: User,
        resource: Resource,
        required_permission: Permission
    ) -> bool:
        # 检查用户是否有足够的权限
        if resource not in user.permissions:
            return False
            
        user_permission = user.permissions[resource]
        return self._is_sufficient_permission(
            user_permission,
            required_permission
        )

数据保护实现

安全框架有了,接下来是具体的数据保护措施:

class DataProtector:
    def __init__(self):
        self.sensitive_patterns = [
            r'\b\d{18}\b',  # 身份证号
            r'\b\d{16}\b',  # 银行卡号
            r'\b\d{11}\b',  # 手机号
            r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'  # 邮箱
        ]
        
    def mask_sensitive_data(
        self,
        text: str,
        security_level: SecurityLevel
    ) -> str:
        if security_level >= SecurityLevel.CONFIDENTIAL:
            # 对高敏感度数据进行掩码
            for pattern in self.sensitive_patterns:
                text = self._apply_mask(text, pattern)
        return text
        
    def _apply_mask(
        self,
        text: str,
        pattern: str
    ) -> str:
        import re
        return re.sub(
            pattern,
            lambda m: '*' * len(m.group()),
            text
        )

class DataEncryptor:
    def __init__(self, key: bytes):
        self.key = key
        self.cipher = AES.new(
            key,
            AES.MODE_GCM
        )
        
    def encrypt(
        self,
        data: str
    ) -> Tuple[bytes, bytes]:
        # 加密数据
        cipher_text, tag = self.cipher.encrypt_and_digest(
            data.encode()
        )
        return cipher_text, tag
        
    def decrypt(
        self,
        cipher_text: bytes,
        tag: bytes
    ) -> str:
        # 解密数据
        plain_text = self.cipher.decrypt_and_verify(
            cipher_text,
            tag
        )
        return plain_text.decode()

行为约束系统

为了防止 AI Agent 执行危险操作,我们需要一个行为约束系统:

class ActionConstraint(BaseModel):
    action: str
    allowed_roles: List[str]
    required_level: SecurityLevel
    max_frequency: Optional[int]
    cooldown: Optional[int]
    
class BehaviorController:
    def __init__(self):
        self.constraints: Dict[str, ActionConstraint] = {}
        self.action_history: Dict[str, List[datetime]] = {}
        
    def add_constraint(
        self,
        constraint: ActionConstraint
    ):
        self.constraints[constraint.action] = constraint
        
    async def check_action(
        self,
        action: str,
        context: SecurityContext
    ) -> bool:
        # 1. 检查是否有约束
        constraint = self.constraints.get(action)
        if not constraint:
            return True
            
        # 2. 检查角色
        if not any(
            role in constraint.allowed_roles
            for role in context.user.roles
        ):
            return False
            
        # 3. 检查安全等级
        if context.user.security_level < constraint.required_level:
            return False
            
        # 4. 检查频率限制
        if not self._check_frequency(
            action,
            constraint,
            context.user.id
        ):
            return False
            
        return True
        
    def _check_frequency(
        self,
        action: str,
        constraint: ActionConstraint,
        user_id: str
    ) -> bool:
        if not constraint.max_frequency:
            return True
            
        key = f"{action}:{user_id}"
        history = self.action_history.get(key, [])
        
        # 清理过期记录
        now = datetime.now()
        history = [
            t for t in history
            if (now - t).total_seconds() <= constraint.cooldown
        ]
        
        # 检查频率
        if len(history) >= constraint.max_frequency:
            return False
            
        # 更新历史
        history.append(now)
        self.action_history[key] = history
        
        return True

审计日志系统

最后,我们需要一个完善的审计系统来记录所有关键操作:

class AuditEvent(BaseModel):
    id: str
    timestamp: datetime
    user_id: str
    action: str
    resource: str
    status: str
    details: Dict[str, Any]
    
class AuditLogger:
    def __init__(self, storage):
        self.storage = storage
        
    async def log_event(
        self,
        context: SecurityContext,
        action: str,
        resource: str,
        status: str,
        details: Dict[str, Any] = None
    ):
        event = AuditEvent(
            id=str(uuid.uuid4()),
            timestamp=datetime.now(),
            user_id=context.user.id,
            action=action,
            resource=resource,
            status=status,
            details=details or {}
        )
        
        # 保存审计日志
        await self.storage.save_event(event)
        
    async def query_events(
        self,
        filters: Dict[str, Any],
        start_time: datetime,
        end_time: datetime,
        limit: int = 100
    ) -> List[AuditEvent]:
        return await self.storage.query_events(
            filters,
            start_time,
            end_time,
            limit
        )

实际应用

把这些组件组合起来,用在我们的 AI Agent 系统中:

class SecureAgent:
    def __init__(
        self,
        security_manager: SecurityManager,
        data_protector: DataProtector,
        behavior_controller: BehaviorController,
        audit_logger: AuditLogger
    ):
        self.security = security_manager
        self.protector = data_protector
        self.controller = behavior_controller
        self.logger = audit_logger
        
    async def process_request(
        self,
        request: Dict[str, Any],
        context: SecurityContext
    ) -> Dict[str, Any]:
        try:
            # 1. 权限检查
            if not await self.security.authorize(
                context.user,
                request["resource"],
                request["permission"]
            ):
                raise PermissionError("权限不足")
                
            # 2. 行为检查
            if not await self.controller.check_action(
                request["action"],
                context
            ):
                raise ValueError("操作不允许")
                
            # 3. 处理请求
            result = await self._handle_request(
                request,
                context
            )
            
            # 4. 数据保护
            protected_result = self.protector.mask_sensitive_data(
                result,
                context.user.security_level
            )
            
            # 5. 记录审计日志
            await self.logger.log_event(
                context,
                request["action"],
                request["resource"],
                "success",
                {"request": request, "result": "masked"}
            )
            
            return protected_result
            
        except Exception as e:
            # 记录错误
            await self.logger.log_event(
                context,
                request["action"],
                request["resource"],
                "error",
                {"error": str(e)}
            )
            raise

最佳实践

在实施这个安全框架的过程中,我总结了几点经验:

  1. 纵深防御

    • 多层次的安全检查
    • 最小权限原则
    • 及时更新和维护
  2. 平衡安全与体验

    • 合理的权限粒度
    • 清晰的错误提示
    • 适度的限制策略
  3. 持续改进

    • 定期安全审计
    • 及时处理反馈
    • 更新安全策略

写在最后

安全不是一个可以一劳永逸的问题,而是需要持续关注和改进的过程。一个好的安全系统应该像一个尽职的保安,既要保护好资产,又不能影响正常工作。

在下一篇文章中,我会讲解如何实现 AI Agent 的测试和监控系统。如果你对安全实践有什么想法,欢迎在评论区交流。


远洋录
3 声望0 粉丝

🚀 独立开发者 | 技术出海实践者