taiji-AI-PAD/test_cre_admins.py

#!/usr/bin/env python3
"""
创建管理员账户脚本
用于创建4个管理员角色：
- 超级管理员 (super_admin)
- 计费管理员 (billing_admin)
- 运维管理员 (operations_admin)
- 渠道管理员 (channel_admin)
"""

import requests
import sys
import os
import asyncio
from typing import Optional
import bcrypt

# 添加services/mcp-server到路径，以便导入模块
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'services', 'mcp-server'))

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy import select
from models import User, Channel, ModelPricing, QuotaAlert
from config import settings

def get_password_hash(password: str) -> str:
    """加密密码（使用bcrypt）"""
    password_bytes = password.encode('utf-8')
    salt = bcrypt.gensalt()
    hashed = bcrypt.hashpw(password_bytes, salt)
    return hashed.decode('utf-8')

BASE_URL = "http://localhost:8002"

# 要创建的管理员列表
ADMINS = [
    {
        "name": "超级管理员",
        "email": "superadmin@taiji-ai.com",
        "password": "Admin@123456",
        "role": "super_admin",
    },
    {
        "name": "计费管理员",
        "email": "newbilling@test.com",
        "password": "Billing@123456",
        "role": "billing_admin",
        "channel_name": "测试渠道",  # 计费管理员也需要关联渠道
        "channel_email": "test-channel@test.com",  # 共享渠道邮箱
    },
    {
        "name": "运维管理员",
        "email": "newops@test.com",
        "password": "Ops@123456",
        "role": "operations_admin",
        "channel_name": "测试渠道",  # 运维管理员也需要关联渠道
        "channel_email": "test-channel@test.com",  # 共享渠道邮箱（与计费管理员共享）
    },
    {
        "name": "渠道管理员",
        "email": "channel-a@test.com",
        "password": "ChannelA@123456",
        "role": "channel_admin",
        "channel_name": "渠道A",  # 渠道名称
    }
]


def login_admin(email: str, password: str, role: str = "super_admin") -> Optional[str]:
    """登录管理员账户，返回token"""
    try:
        resp = requests.post(
            f"{BASE_URL}/api/auth/login",
            json={
                "email": email,
                "password": password,
                "role": role
            },
            timeout=10
        )
        
        if resp.status_code == 200:
            data = resp.json()
            token = data.get("data", {}).get("token")
            if token:
                print(f"  ✓ 登录成功: {email}")
                return token
            else:
                print(f"  ✗ 登录失败: 响应中未找到token")
                return None
        else:
            error = resp.json().get("detail", resp.text)
            print(f"  ✗ 登录失败: {error}")
            return None
    except Exception as e:
        print(f"  ✗ 登录出错: {e}")
        return None


async def create_admin_via_api(token: str, admin_info: dict) -> bool:
    """通过API创建管理员（需要先创建渠道）"""
    try:
        # 注意：API只能创建billing_admin和operations_admin
        # 超级管理员和渠道管理员需要直接操作数据库
        if admin_info["role"] not in ["billing_admin", "operations_admin"]:
            print(f"  ⚠ 跳过: {admin_info['role']} 需要通过数据库直接创建")
            return False
        
        # 先创建或获取渠道
        channel_id = None
        if admin_info.get("channel_name"):
            channel_id = await get_or_create_channel_for_api(admin_info)
        
        # 构建请求数据
        request_data = {
            "name": admin_info["name"],
            "email": admin_info["email"],
            "password": admin_info["password"],
            "role": admin_info["role"]
        }
        if channel_id:
            request_data["channelId"] = str(channel_id)
            
        resp = requests.post(
            f"{BASE_URL}/api/admin/admins/create",
            headers={
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json"
            },
            json=request_data,
            timeout=10
        )
        
        if resp.status_code == 200:
            data = resp.json()
            channel_id_from_response = data.get("data", {}).get("channelId")
            print(f"  ✓ 创建成功: {admin_info['email']}" + (f" (渠道ID: {channel_id_from_response})" if channel_id_from_response else ""))
            return True
        else:
            error = resp.json().get("detail", resp.text)
            if "邮箱已被使用" in error or "already exists" in error.lower():
                print(f"  ⚠ 已存在: {admin_info['email']}")
                # 如果已存在，尝试更新channel_id
                if channel_id:
                    await update_existing_user_channel(admin_info["email"], channel_id)
                return True  # 已存在也算成功
            else:
                print(f"  ✗ 创建失败: {error}")
                return False
    except Exception as e:
        print(f"  ✗ 创建出错: {e}")
        import traceback
        traceback.print_exc()
        return False


async def get_or_create_channel_for_api(admin_info: dict) -> str:
    """为API创建获取或创建渠道，返回channel_id字符串"""
    try:
        database_url = settings.database_url
        if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        
        engine = create_async_engine(database_url, echo=False)
        AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        
        async with AsyncSessionLocal() as session:
            channel_name = admin_info.get("channel_name", f"渠道-{admin_info['name']}")
            channel_email = admin_info.get("channel_email", f"channel-{channel_name.lower().replace(' ', '-')}@test.com")
            channel = await get_or_create_channel(session, channel_email, channel_name)
            await session.commit()
            channel_id = str(channel.id)
            await engine.dispose()
            return channel_id
    except Exception as e:
        print(f"  ⚠ 创建渠道失败: {e}")
        return None


async def update_existing_user_channel(email: str, channel_id: str):
    """更新已存在用户的channel_id"""
    try:
        database_url = settings.database_url
        if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        
        engine = create_async_engine(database_url, echo=False)
        AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        
        async with AsyncSessionLocal() as session:
            result = await session.execute(select(User).where(User.email == email))
            user = result.scalar_one_or_none()
            if user:
                import uuid
                user.channel_id = uuid.UUID(channel_id)
                await session.commit()
                print(f"  ✓ 已更新用户的渠道ID: {channel_id}")
        await engine.dispose()
    except Exception as e:
        print(f"  ⚠ 更新用户渠道ID失败: {e}")


async def verify_database_migration():
    """验证数据库迁移是否成功"""
    engine = None
    try:
        database_url = settings.database_url
        if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        
        engine = create_async_engine(database_url, echo=False)
        AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        
        async with AsyncSessionLocal() as session:
            # 检查model_pricing表
            try:
                result = await session.execute(select(ModelPricing).limit(1))
                print("  ✓ model_pricing 表存在")
            except Exception as e:
                print(f"  ✗ model_pricing 表不存在或有问题: {e}")
            
            # 检查quota_alerts表
            try:
                result = await session.execute(select(QuotaAlert).limit(1))
                print("  ✓ quota_alerts 表存在")
            except Exception as e:
                print(f"  ✗ quota_alerts 表不存在或有问题: {e}")
        
    except Exception as e:
        print(f"  ✗ 数据库连接失败: {e}")
    finally:
        if engine:
            await engine.dispose()


async def get_or_create_channel(session: AsyncSession, channel_email: str, channel_name: str) -> Channel:
    """获取或创建渠道"""
    # 先查找是否已存在
    result = await session.execute(
        select(Channel).where(Channel.email == channel_email)
    )
    channel = result.scalar_one_or_none()
    
    if channel:
        return channel
    
    # 创建新渠道
    channel = Channel(
        name=channel_name,
        email=channel_email,
        password_hash=get_password_hash("Channel@123456"),  # 默认密码
        commission_rate=10.0,
        channel_credit=0,
        custom_agent_cpu=2,
        custom_agent_memory=4,
        status="active",
    )
    session.add(channel)
    await session.flush()  # 获取ID但不提交
    await session.refresh(channel)
    print(f"  ✓ 创建渠道: {channel_name} (ID: {channel.id})")
    return channel


async def create_admin_via_db(admin_info: dict) -> bool:
    """直接通过数据库创建管理员"""
    engine = None
    try:
        # 准备数据库URL
        database_url = settings.database_url
        if database_url.startswith("postgresql://") and "+asyncpg" not in database_url:
            database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
        
        # 创建数据库引擎
        engine = create_async_engine(database_url, echo=False)
        AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        
        async with AsyncSessionLocal() as session:
            # 如果是渠道管理员、计费管理员或运维管理员，需要先创建或获取渠道
            channel_id = None
            if admin_info["role"] in ["channel_admin", "billing_admin", "operations_admin"]:
                # 为管理员创建对应的渠道
                channel_name = admin_info.get("channel_name", f"渠道-{admin_info['name']}")
                # 使用一个统一的渠道邮箱（如果多个管理员共享同一个渠道）
                channel_email = admin_info.get("channel_email", f"channel-{channel_name.lower().replace(' ', '-')}@test.com")
                channel = await get_or_create_channel(session, channel_email, channel_name)
                channel_id = channel.id
                await session.commit()  # 提交渠道创建
            
            # 检查用户是否已存在
            result = await session.execute(
                select(User).where(User.email == admin_info["email"])
            )
            existing_user = result.scalar_one_or_none()
            
            if existing_user:
                print(f"  ⚠ 用户已存在: {admin_info['email']}")
                # 更新角色和密码
                existing_user.role = admin_info["role"]
                existing_user.password_hash = get_password_hash(admin_info["password"])
                existing_user.hashed_password = existing_user.password_hash
                existing_user.name = admin_info["name"]
                existing_user.username = admin_info["email"].split("@")[0]
                existing_user.full_name = admin_info["name"]
                existing_user.is_active = True
                if admin_info["role"] == "super_admin":
                    existing_user.is_admin = True
                # 如果是需要渠道的角色，更新channel_id
                if channel_id and admin_info["role"] in ["channel_admin", "billing_admin", "operations_admin"]:
                    existing_user.channel_id = channel_id
                await session.commit()
                print(f"  ✓ 更新成功: {admin_info['email']}" + (f" (渠道ID: {channel_id})" if channel_id else ""))
                return True
            
            # 创建新用户
            password_hash = get_password_hash(admin_info["password"])
            user = User(
                name=admin_info["name"],
                email=admin_info["email"],
                password_hash=password_hash,
                hashed_password=password_hash,
                username=admin_info["email"].split("@")[0],
                full_name=admin_info["name"],
                role=admin_info["role"],
                channel_id=channel_id,  # 关联渠道ID
                is_active=True,
                is_admin=(admin_info["role"] == "super_admin"),
                status="active",
                balance=0,
                credit_limit=0,
                # 新增字段支持
                eu_balance=0,  # EU余额
                total_eu_consumed=0,  # 总EU消耗
            )
            
            session.add(user)
            await session.commit()
            await session.refresh(user)
            
            print(f"  ✓ 创建成功: {admin_info['email']} (角色: {admin_info['role']})" + (f" (渠道ID: {channel_id})" if channel_id else ""))
            return True
            
    except Exception as e:
        print(f"  ✗ 数据库创建失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        if engine:
            await engine.dispose()


async def main():
    """主函数"""
    print("="*80)
    print("创建管理员账户")
    print("="*80)
    print()
    
    # 验证新表是否创建成功
    print("步骤0: 验证数据库迁移...")
    await verify_database_migration()
    print()
    
    # 首先尝试登录默认admin账户
    print("步骤1: 尝试登录默认管理员账户...")
    default_admin_email = "admin@taiji-ai.com"
    default_admin_password = "admin123"
    
    token = login_admin(default_admin_email, default_admin_password, "super_admin")
    
    # 如果没有默认admin，尝试创建超级管理员
    if not token:
        print("\n步骤2: 默认管理员不存在，直接创建超级管理员...")
        super_admin = ADMINS[0]  # 第一个是超级管理员
        success = await create_admin_via_db(super_admin)
        if success:
            print("\n步骤3: 使用新创建的超级管理员登录...")
            token = login_admin(super_admin["email"], super_admin["password"], "super_admin")
        else:
            print("  ✗ 无法创建超级管理员，请检查数据库连接")
            return
    
    if not token:
        print("  ✗ 无法获取管理员token，请检查服务是否运行")
        return
    
    print(f"\n步骤4: 创建其他管理员账户...")
    print("-" * 80)
    
    results = []
    for admin in ADMINS:
        print(f"\n创建 {admin['name']} ({admin['email']})...")
        
        # 超级管理员和渠道管理员需要直接操作数据库
        if admin["role"] in ["super_admin", "channel_admin"]:
            success = await create_admin_via_db(admin)
        else:
            # billing_admin和operations_admin可以通过API创建
            success = await create_admin_via_api(token, admin)
        
        results.append({
            "name": admin["name"],
            "email": admin["email"],
            "role": admin["role"],
            "success": success
        })
    
    # 输出结果汇总
    print("\n" + "="*80)
    print("创建结果汇总")
    print("="*80)
    print(f"\n{'角色':<20} {'邮箱':<35} {'状态'}")
    print("-" * 80)
    
    for result in results:
        status = "✓ 成功" if result["success"] else "✗ 失败"
        print(f"{result['name']:<20} {result['email']:<35} {status}")
    
    success_count = sum(1 for r in results if r["success"])
    print(f"\n总计: {success_count}/{len(results)} 个账户创建成功")
    
    # 输出账户信息
    print("\n" + "="*80)
    print("账户信息")
    print("="*80)
    for admin in ADMINS:
        print(f"{admin['name']:<20} | {admin['email']:<35} | 密码: {admin['password']}")


if __name__ == "__main__":
    asyncio.run(main())