aiagent/backend/scripts/init_rbac_data.py

#!/usr/bin/env python3
"""
初始化RBAC数据
创建系统角色和权限
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.database import SessionLocal
from app.models.permission import Role, Permission
import uuid

# 系统角色定义
SYSTEM_ROLES = [
    {
        "name": "admin",
        "description": "系统管理员，拥有所有权限",
        "is_system": True
    },
    {
        "name": "developer",
        "description": "开发者，可以创建和管理工作流、Agent",
        "is_system": True
    },
    {
        "name": "viewer",
        "description": "查看者，只能查看工作流和执行记录",
        "is_system": True
    },
    {
        "name": "operator",
        "description": "操作员，可以执行工作流，但不能修改",
        "is_system": True
    }
]

# 权限定义
PERMISSIONS = [
    # 工作流权限
    {"name": "工作流-创建", "code": "workflow:create", "resource": "workflow", "action": "create", "description": "创建工作流"},
    {"name": "工作流-查看", "code": "workflow:read", "resource": "workflow", "action": "read", "description": "查看工作流"},
    {"name": "工作流-更新", "code": "workflow:update", "resource": "workflow", "action": "update", "description": "更新工作流"},
    {"name": "工作流-删除", "code": "workflow:delete", "resource": "workflow", "action": "delete", "description": "删除工作流"},
    {"name": "工作流-执行", "code": "workflow:execute", "resource": "workflow", "action": "execute", "description": "执行工作流"},
    {"name": "工作流-分享", "code": "workflow:share", "resource": "workflow", "action": "share", "description": "分享工作流"},
    
    # Agent权限
    {"name": "Agent-创建", "code": "agent:create", "resource": "agent", "action": "create", "description": "创建Agent"},
    {"name": "Agent-查看", "code": "agent:read", "resource": "agent", "action": "read", "description": "查看Agent"},
    {"name": "Agent-更新", "code": "agent:update", "resource": "agent", "action": "update", "description": "更新Agent"},
    {"name": "Agent-删除", "code": "agent:delete", "resource": "agent", "action": "delete", "description": "删除Agent"},
    {"name": "Agent-执行", "code": "agent:execute", "resource": "agent", "action": "execute", "description": "执行Agent"},
    {"name": "Agent-部署", "code": "agent:deploy", "resource": "agent", "action": "deploy", "description": "部署Agent"},
    
    # 执行权限
    {"name": "执行-查看", "code": "execution:read", "resource": "execution", "action": "read", "description": "查看执行记录"},
    {"name": "执行-取消", "code": "execution:cancel", "resource": "execution", "action": "cancel", "description": "取消执行"},
    
    # 数据源权限
    {"name": "数据源-创建", "code": "data_source:create", "resource": "data_source", "action": "create", "description": "创建数据源"},
    {"name": "数据源-查看", "code": "data_source:read", "resource": "data_source", "action": "read", "description": "查看数据源"},
    {"name": "数据源-更新", "code": "data_source:update", "resource": "data_source", "action": "update", "description": "更新数据源"},
    {"name": "数据源-删除", "code": "data_source:delete", "resource": "data_source", "action": "delete", "description": "删除数据源"},
    
    # 模型配置权限
    {"name": "模型配置-创建", "code": "model_config:create", "resource": "model_config", "action": "create", "description": "创建模型配置"},
    {"name": "模型配置-查看", "code": "model_config:read", "resource": "model_config", "action": "read", "description": "查看模型配置"},
    {"name": "模型配置-更新", "code": "model_config:update", "resource": "model_config", "action": "update", "description": "更新模型配置"},
    {"name": "模型配置-删除", "code": "model_config:delete", "resource": "model_config", "action": "delete", "description": "删除模型配置"},
    
    # 权限管理权限
    {"name": "权限-管理", "code": "permission:manage", "resource": "permission", "action": "manage", "description": "管理权限和角色"},
]

# 角色权限映射
ROLE_PERMISSIONS = {
    "admin": ["*"],  # 所有权限
    "developer": [
        "workflow:create", "workflow:read", "workflow:update", "workflow:delete", "workflow:execute", "workflow:share",
        "agent:create", "agent:read", "agent:update", "agent:delete", "agent:execute", "agent:deploy",
        "execution:read", "execution:cancel",
        "data_source:create", "data_source:read", "data_source:update", "data_source:delete",
        "model_config:create", "model_config:read", "model_config:update", "model_config:delete"
    ],
    "viewer": [
        "workflow:read",
        "agent:read",
        "execution:read",
        "data_source:read",
        "model_config:read"
    ],
    "operator": [
        "workflow:read", "workflow:execute",
        "agent:read", "agent:execute",
        "execution:read", "execution:cancel"
    ]
}


def init_rbac_data():
    """初始化RBAC数据"""
    db = SessionLocal()
    try:
        print("=" * 60)
        print("初始化RBAC数据")
        print("=" * 60)
        print()
        
        # 创建权限
        print("创建权限...")
        permission_map = {}
        for perm_data in PERMISSIONS:
            existing = db.query(Permission).filter(Permission.code == perm_data["code"]).first()
            if existing:
                print(f"  权限已存在: {perm_data['code']}")
                permission_map[perm_data["code"]] = existing
            else:
                permission = Permission(
                    id=str(uuid.uuid4()),
                    name=perm_data["name"],
                    code=perm_data["code"],
                    resource=perm_data["resource"],
                    action=perm_data["action"],
                    description=perm_data["description"]
                )
                db.add(permission)
                permission_map[perm_data["code"]] = permission
                print(f"  ✅ 创建权限: {perm_data['code']}")
        
        db.commit()
        print()
        
        # 创建角色
        print("创建角色...")
        role_map = {}
        for role_data in SYSTEM_ROLES:
            existing = db.query(Role).filter(Role.name == role_data["name"]).first()
            if existing:
                print(f"  角色已存在: {role_data['name']}")
                role_map[role_data["name"]] = existing
            else:
                role = Role(
                    id=str(uuid.uuid4()),
                    name=role_data["name"],
                    description=role_data["description"],
                    is_system=role_data["is_system"]
                )
                db.add(role)
                role_map[role_data["name"]] = role
                print(f"  ✅ 创建角色: {role_data['name']}")
        
        db.commit()
        print()
        
        # 分配权限给角色
        print("分配权限给角色...")
        for role_name, permission_codes in ROLE_PERMISSIONS.items():
            role = role_map.get(role_name)
            if not role:
                continue
            
            if permission_codes == ["*"]:
                # 管理员拥有所有权限
                role.permissions = list(permission_map.values())
                print(f"  ✅ {role_name}: 分配所有权限")
            else:
                # 分配指定权限
                permissions = [permission_map[code] for code in permission_codes if code in permission_map]
                role.permissions = permissions
                print(f"  ✅ {role_name}: 分配 {len(permissions)} 个权限")
        
        db.commit()
        print()
        
        print("=" * 60)
        print("✅ RBAC数据初始化完成！")
        print("=" * 60)
        
    except Exception as e:
        db.rollback()
        print(f"❌ 初始化失败: {e}")
        import traceback
        traceback.print_exc()
    finally:
        db.close()


if __name__ == "__main__":
    init_rbac_data()