Ubuntu
199 lines
5.7 KiB
Python
199 lines
5.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
安全控制模块
|
|
负责权限检查、内容过滤、速率限制
|
|
"""
|
|
|
|
import re
|
|
import time
|
|
import logging
|
|
from typing import Dict, List, Set, Tuple
|
|
from collections import defaultdict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SecurityChecker:
|
|
"""安全检查器"""
|
|
|
|
def __init__(self, config: Dict):
|
|
self.config = config
|
|
|
|
# 允许的工具
|
|
self.allowed_tools: Set[str] = set(
|
|
config.get('allowed_tools', ['read', 'write', 'edit', 'exec'])
|
|
)
|
|
|
|
# 禁止的命令
|
|
self.blocked_commands: List[str] = config.get(
|
|
'blocked_commands',
|
|
['rm -rf /', 'dd if=', 'mkfs', 'shutdown', 'reboot', 'halt']
|
|
)
|
|
|
|
# 敏感关键词
|
|
self.sensitive_keywords: List[str] = [
|
|
'password', 'secret', 'token', 'api_key',
|
|
'私钥', '密码', '密钥', '.ssh/id_rsa'
|
|
]
|
|
|
|
# 速率限制
|
|
self.max_requests = config.get('max_requests_per_minute', 10)
|
|
self.request_records: Dict[str, List[float]] = defaultdict(list)
|
|
|
|
def check_rate_limit(self, user_id: str) -> bool:
|
|
"""
|
|
检查速率限制
|
|
|
|
Args:
|
|
user_id: 用户标识
|
|
|
|
Returns:
|
|
是否允许请求
|
|
"""
|
|
now = time.time()
|
|
window = 60 # 60秒窗口
|
|
|
|
# 清理过期记录
|
|
self.request_records[user_id] = [
|
|
t for t in self.request_records[user_id]
|
|
if now - t < window
|
|
]
|
|
|
|
# 检查是否超过限制
|
|
if len(self.request_records[user_id]) >= self.max_requests:
|
|
logger.warning(f"用户 {user_id} 触发速率限制")
|
|
return False
|
|
|
|
# 记录本次请求
|
|
self.request_records[user_id].append(now)
|
|
return True
|
|
|
|
def check_content(self, content: str) -> bool:
|
|
"""
|
|
检查内容是否合法
|
|
|
|
Args:
|
|
content: 消息内容
|
|
|
|
Returns:
|
|
是否合法
|
|
"""
|
|
# 检查敏感关键词
|
|
content_lower = content.lower()
|
|
for keyword in self.sensitive_keywords:
|
|
if keyword.lower() in content_lower:
|
|
logger.warning(f"内容包含敏感词: {keyword}")
|
|
# 不直接拒绝，只记录警告
|
|
# 实际使用时可根据需要调整策略
|
|
|
|
# 检查是否有明显的恶意命令
|
|
for cmd in self.blocked_commands:
|
|
if cmd in content:
|
|
logger.warning(f"检测到禁止命令: {cmd}")
|
|
return False
|
|
|
|
return True
|
|
|
|
def check_tool_permission(self, tool_name: str) -> bool:
|
|
"""
|
|
检查工具使用权限
|
|
|
|
Args:
|
|
tool_name: 工具名称
|
|
|
|
Returns:
|
|
是否允许使用
|
|
"""
|
|
return tool_name in self.allowed_tools
|
|
|
|
def check_command_safety(self, command: str) -> Tuple[bool, str]:
|
|
"""
|
|
检查命令是否安全
|
|
|
|
Args:
|
|
command: 命令字符串
|
|
|
|
Returns:
|
|
(是否安全, 原因)
|
|
"""
|
|
command_lower = command.lower().strip()
|
|
|
|
# 检查危险命令模式
|
|
dangerous_patterns = [
|
|
(r'^\s*rm\s+-rf\s*/', "禁止删除根目录"),
|
|
(r'^\s*dd\s+if=', "禁止直接使用 dd 命令"),
|
|
(r'^\s*mkfs', "禁止格式化文件系统"),
|
|
(r'^\s*shutdown', "禁止关机命令"),
|
|
(r'^\s*reboot', "禁止重启命令"),
|
|
(r'^\s*halt', "禁止 halt 命令"),
|
|
(r'>\s*/dev/sd[a-z]', "禁止直接操作磁盘设备"),
|
|
(r'>\s*/dev/null.*<\s*/dev/random', "禁止危险的重定向操作"),
|
|
]
|
|
|
|
for pattern, reason in dangerous_patterns:
|
|
if re.search(pattern, command_lower):
|
|
return False, reason
|
|
|
|
# 检查敏感文件访问
|
|
sensitive_paths = [
|
|
'/etc/passwd',
|
|
'/etc/shadow',
|
|
'~/.ssh/id_rsa',
|
|
'~/.aws/credentials',
|
|
]
|
|
|
|
for path in sensitive_paths:
|
|
if path in command:
|
|
return False, f"禁止访问敏感文件: {path}"
|
|
|
|
return True, "安全"
|
|
|
|
def sanitize_output(self, output: str) -> str:
|
|
"""
|
|
清理输出内容
|
|
|
|
Args:
|
|
output: 原始输出
|
|
|
|
Returns:
|
|
清理后的输出
|
|
"""
|
|
# 移除可能的敏感信息
|
|
patterns = [
|
|
(r'[a-zA-Z0-9_-]{20,}\.[a-zA-Z0-9_-]{20,}', '[TOKEN_REDACTED]'),
|
|
(r'sk-[a-zA-Z0-9]{20,}', '[API_KEY_REDACTED]'),
|
|
(r'password\s*[:=]\s*\S+', 'password: [REDACTED]'),
|
|
]
|
|
|
|
result = output
|
|
for pattern, replacement in patterns:
|
|
result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)
|
|
|
|
return result
|
|
|
|
def get_blocked_commands_list(self) -> List[str]:
|
|
"""获取禁止命令列表"""
|
|
return self.blocked_commands.copy()
|
|
|
|
|
|
class SimpleSecurityChecker:
|
|
"""简化版安全检查器（用于测试）"""
|
|
|
|
def __init__(self, config: Dict):
|
|
pass
|
|
|
|
def check_rate_limit(self, user_id: str) -> bool:
|
|
return True
|
|
|
|
def check_content(self, content: str) -> bool:
|
|
return True
|
|
|
|
def check_tool_permission(self, tool_name: str) -> bool:
|
|
return True
|
|
|
|
def check_command_safety(self, command: str) -> Tuple[bool, str]:
|
|
return True, "安全"
|
|
|
|
def sanitize_output(self, output: str) -> str:
|
|
return output
|