AI Agent 权限管理与访问控制最佳实践

from enum import Enum
from dataclasses import dataclass, field
from typing import Set, Dict, Optional
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)

class PermissionLevel(Enum):
    READ = "read"
    WRITE = "write"
    EXECUTE = "execute"
    DELETE = "delete"
    ADMIN = "admin"

class ResourceType(Enum):
    FILE_SYSTEM = "file_system"
    DATABASE = "database"
    API_ENDPOINT = "api_endpoint"
    NETWORK = "network"
    ENVIRONMENT = "environment"
    EXTERNAL_SERVICE = "external_service"

@dataclass
class Permission:
    """权限定义"""
    resource_type: ResourceType
    level: PermissionLevel
    resource_pattern: str  # 支持通配符匹配
    
    def matches(self, req_type: ResourceType, req_level: PermissionLevel, 
                req_resource: str) -> bool:
        """检查权限是否匹配请求"""
        if self.resource_type != req_type:
            return False
        if self.level != req_level and self.level != PermissionLevel.ADMIN:
            return False
        # 简单的通配符匹配
        import fnmatch
        return fnmatch.fnmatch(req_resource, self.resource_pattern)

@dataclass
class AgentProfile:
    """Agent 权限档案"""
    agent_id: str
    role: str
    permissions: Set[Permission] = field(default_factory=set)
    temporary_permissions: Dict[Permission, datetime] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    trust_level: int = 1  # 1-5，越高越信任
    violation_count: int = 0
    
    def has_permission(self, req_type: ResourceType, req_level: PermissionLevel,
                       req_resource: str) -> bool:
        """检查 Agent 是否有指定权限"""
        # 检查永久权限
        for perm in self.permissions:
            if perm.matches(req_type, req_level, req_resource):
                return True
        
        # 检查临时权限（检查是否过期）
        now = datetime.now()
        expired = []
        for perm, expiry in self.temporary_permissions.items():
            if now < expiry:
                if perm.matches(req_type, req_level, req_resource):
                    return True
            else:
                expired.append(perm)
        
        # 清理过期权限
        for perm in expired:
            del self.temporary_permissions[perm]
        
        return False

class PermissionGuard:
    """权限守卫：执行最小权限检查"""
    
    def __init__(self):
        self._agents: Dict[str, AgentProfile] = {}
        self._default_deny = True  # 默认拒绝所有未明确允许的操作
    
    def register_agent(self, agent_id: str, role: str, 
                       permissions: Set[Permission],
                       trust_level: int = 1) -> AgentProfile:
        """注册 Agent 并授予最小权限集"""
        profile = AgentProfile(
            agent_id=agent_id,
            role=role,
            permissions=permissions,
            trust_level=trust_level
        )
        self._agents[agent_id] = profile
        logger.info(f"Agent {agent_id} registered with role={role}, "
                    f"permissions={len(permissions)}")
        return profile
    
    def grant_temporary(self, agent_id: str, permission: Permission,
                        duration: timedelta) -> bool:
        """授予临时权限（带过期时间）"""
        if agent_id not in self._agents:
            return False
        
        profile = self._agents[agent_id]
        expiry = datetime.now() + duration
        profile.temporary_permissions[permission] = expiry
        logger.info(f"Granted temp permission to {agent_id}: "
                    f"{permission} expires at {expiry}")
        return True
    
    def check_permission(self, agent_id: str, req_type: ResourceType,
                         req_level: PermissionLevel, req_resource: str) -> bool:
        """检查权限，返回是否允许"""
        if agent_id not in self._agents:
            logger.warning(f"Unknown agent: {agent_id}")
            return False
        
        profile = self._agents[agent_id]
        allowed = profile.has_permission(req_type, req_level, req_resource)
        
        if not allowed:
            profile.violation_count += 1
            logger.warning(f"Permission DENIED: agent={agent_id}, "
                          f"type={req_type}, level={req_level}, "
                          f"resource={req_resource}, violations="
                          f"{profile.violation_count}")
            
            # 违规次数过多时自动降级
            if profile.violation_count >= 5:
                profile.trust_level = max(1, profile.trust_level - 1)
                logger.warning(f"Agent {agent_id} trust level reduced to "
                             f"{profile.trust_level}")
        else:
            logger.debug(f"Permission GRANTED: agent={agent_id}, "
                        f"type={req_type}, level={req_level}")
        
        return allowed
    
    def audit_log(self, agent_id: str) -> dict:
        """生成 Agent 的审计报告"""
        if agent_id not in self._agents:
            return {}
        profile = self._agents[agent_id]
        return {
            "agent_id": agent_id,
            "role": profile.role,
            "trust_level": profile.trust_level,
            "violation_count": profile.violation_count,
            "active_permissions": len(profile.permissions),
            "temp_permissions": len(profile.temporary_permissions),
            "registered_at": profile.created_at.isoformat()
        }

# === 使用示例 ===
def setup_data_reader_agent(guard: PermissionGuard):
    """示例：创建一个只读数据 Agent"""
    permissions = {
        Permission(ResourceType.DATABASE, PermissionLevel.READ, "analytics_*"),
        Permission(ResourceType.DATABASE, PermissionLevel.READ, "reports_*"),
        Permission(ResourceType.FILE_SYSTEM, PermissionLevel.READ, "/data/reports/*"),
        Permission(ResourceType.API_ENDPOINT, PermissionLevel.READ, "/api/v1/query"),
    }
    return guard.register_agent(
        agent_id="data-reader-001",
        role="data_reader",
        permissions=permissions,
        trust_level=3
    )

def setup_file_cleanup_agent(guard: PermissionGuard):
    """示例：创建文件清理 Agent（仅允许删除特定目录的临时文件）"""
    permissions = {
        Permission(ResourceType.FILE_SYSTEM, PermissionLevel.READ, "/tmp/agent-workspace/*"),
        Permission(ResourceType.FILE_SYSTEM, PermissionLevel.DELETE, "/tmp/agent-workspace/*.tmp"),
        Permission(ResourceType.FILE_SYSTEM, PermissionLevel.WRITE, "/tmp/agent-workspace/logs/"),
    }
    return guard.register_agent(
        agent_id="file-cleanup-001",
        role="file_cleanup",
        permissions=permissions,
        trust_level=2
    )

# 测试
guard = PermissionGuard()
reader = setup_data_reader_agent(guard)
cleaner = setup_file_cleanup_agent(guard)

# 允许的操作
assert guard.check_permission("data-reader-001", ResourceType.DATABASE,
                              PermissionLevel.READ, "analytics_users")
assert guard.check_permission("file-cleanup-001", ResourceType.FILE_SYSTEM,
                              PermissionLevel.DELETE, "/tmp/agent-workspace/cache.tmp")

# 拒绝的操作
assert not guard.check_permission("data-reader-001", ResourceType.DATABASE,
                                   PermissionLevel.DELETE, "analytics_users")
assert not guard.check_permission("file-cleanup-001", ResourceType.DATABASE,
                                   PermissionLevel.READ, "production_db")

print("所有权限检查通过！最小权限守卫系统正常工作。")

from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from enum import Enum
import re

class PolicyEffect(Enum):
    ALLOW = "allow"
    DENY = "deny"

@dataclass
class PolicyRule:
    """ABAC 策略规则"""
    name: str
    effect: PolicyEffect
    conditions: Dict[str, Any]
    priority: int = 0
    
    def evaluate(self, context: Dict[str, Any]) -> bool:
        for key, expected in self.conditions.items():
            actual = self._resolve_key(context, key)
            if actual is None:
                return False
            if not self._match(actual, expected):
                return False
        return True
    
    def _resolve_key(self, context: Dict[str, Any], key: str) -> Any:
        parts = key.split(".")
        current = context
        for part in parts:
            if isinstance(current, dict) and part in current:
                current = current[part]
            else:
                return None
        return current
    
    def _match(self, actual: Any, expected: Any) -> bool:
        if isinstance(expected, str) and expected.startswith("regex:"):
            return bool(re.match(expected[6:], str(actual)))
        if isinstance(expected, str) and expected.startswith("<"):
            return str(actual) < expected[1:]
        if isinstance(expected, str) and expected.startswith(">"):
            return str(actual) > expected[1:]
        if isinstance(expected, list):
            return actual in expected
        return actual == expected

class ABACPolicyEngine:
    def __init__(self):
        self._rules: List[PolicyRule] = []
    
    def add_rule(self, rule: PolicyRule):
        self._rules.append(rule)
        self._rules.sort(key=lambda r: r.priority, reverse=True)
    
    def evaluate(self, context: Dict[str, Any]) -> PolicyEffect:
        for rule in self._rules:
            if rule.evaluate(context):
                return rule.effect
        return PolicyEffect.DENY

# 使用示例
engine = ABACPolicyEngine()
engine.add_rule(PolicyRule(
    name="deny-production-delete", effect=PolicyEffect.DENY,
    conditions={"action.type": "delete", "object.environment": "production"}, priority=100))
engine.add_rule(PolicyRule(
    name="admin-delete-work-hours", effect=PolicyEffect.ALLOW,
    conditions={"subject.role": "admin", "action.type": "delete", "env.time_hour": [9,10,11,12,13,14,15,16,17]}, priority=10))

result = engine.evaluate({"subject": {"role": "admin"}, "action": {"type": "delete"},
    "object": {"environment": "production"}, "env": {"time_hour": 14}})
print(f"决策结果: {result.value}")  # DENY（生产环境最高优先级）