在前面的文章中,我们讨论了 AI Agent 的各个功能模块。今天,我想聊聊一个经常被忽视但极其重要的话题:安全性。说实话,我在这方面也吃过亏,希望通过分享我的经验,帮大家少走一些弯路。
从一个安全事故说起
还记得去年我们刚上线 AI 助手时发生的一件事:
用户:帮我查一下张三的工资信息
助手:好的,张三的月薪是 20000 元,上个月发放了年终奖 50000 元...
用户:...(这是个普通用户,不应该有权限查看这些信息)
这个事故让我意识到:AI Agent 不仅要能完成任务,还要懂得"什么能做,什么不能做"。于是我开始系统性地思考 AI Agent 的安全问题:
- 权限控制:谁能访问什么
- 数据保护:如何保护敏感信息
- 行为约束:如何限制危险操作
- 审计追踪:记录关键操作
安全框架设计
经过多次迭代,我设计了一个相对完善的安全框架:
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
最佳实践
在实施这个安全框架的过程中,我总结了几点经验:
纵深防御
- 多层次的安全检查
- 最小权限原则
- 及时更新和维护
平衡安全与体验
- 合理的权限粒度
- 清晰的错误提示
- 适度的限制策略
持续改进
- 定期安全审计
- 及时处理反馈
- 更新安全策略
写在最后
安全不是一个可以一劳永逸的问题,而是需要持续关注和改进的过程。一个好的安全系统应该像一个尽职的保安,既要保护好资产,又不能影响正常工作。
在下一篇文章中,我会讲解如何实现 AI Agent 的测试和监控系统。如果你对安全实践有什么想法,欢迎在评论区交流。
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用
。你还可以使用@
来通知其他用户。