非你网

运行起 dify 后，在控制台里可以看到登录调用的接口文件定义在这里：api\controllers\console\auth\login.py 现在想扩展新的接口，可以在当前目录下新建一个 register.py 文件，加入以下代码。

import os
import hashlib
import functools
import time
import logging
import secrets
import base64

from flask import make_response, request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator, ValidationError
from sqlalchemy import select, update
from typing import Literal, Optional

from controllers.common.fields import SimpleResultOptionalDataResponse
from controllers.common.schema import (
    register_response_schema_models,
    register_schema_models,
)
from controllers.console import console_ns
from libs.helper import EmailStr
from libs.helper import timezone as validate_timezone_string
from libs.password import hash_password
from constants.languages import get_valid_language
from models import Account, Tenant, TenantAccountJoin
from extensions.ext_database import db
from services.errors.account import AccountRegisterError
from controllers.console.wraps import (
    decrypt_password_field,
    setup_required,
)

from sqlalchemy import func
from models import ApiToken

logger = logging.getLogger(__name__)

# 签名校验装饰器
def private_api_auth_required(f):
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        private_secret = os.environ.get("CONSOLE_PRIVATE_API_SECRET")
        if not private_secret or len(private_secret) != 12:
            return {"result": "fail", "message": "Private API service unavailable"}, 503

        req_time = request.headers.get("X-API-TIME", "")
        req_sign = request.headers.get("X-API-SIGN", "")

        if not req_time or not req_sign:
            return {"result": "fail", "message": "Missing required headers"}, 400

        try:
            timestamp = int(req_time)
            if timestamp > 9999999999:
                timestamp = timestamp // 1000
        except (ValueError, TypeError):
            return {"result": "fail", "message": "Invalid timestamp"}, 400

        if abs(int(time.time()) - timestamp) > 300:
            return {"result": "fail", "message": "Request expired"}, 401

        calc_sign = (
            hashlib.md5(f"{private_secret}{req_time}".encode()).hexdigest().lower()
        )
        if calc_sign != req_sign.lower():
            return {"result": "fail", "message": "Invalid signature"}, 401

        return f(*args, **kwargs)

    return decorated

# 请求模型：字段改为 name，timezone 默认 Asia/Shanghai
class RegisterPayload(BaseModel):
    email: EmailStr = Field(..., description="邮箱")
    name: str = Field(..., min_length=1, max_length=50, description="用户名")
    password: str = Field(..., description="明文密码")
    language: str | None = Field(default=None, description="界面语言")
    timezone: str = Field(
        default="Asia/Shanghai", description="时区，默认 Asia/Shanghai"
    )

    @field_validator("timezone")
    @classmethod
    def validate_timezone(cls, value):
        return validate_timezone_string(value) if value else "Asia/Shanghai"

# 获取用户详情
class UserInfoPayload(BaseModel):
    account_id: str = Field(..., description="用户ID")

class GenerateApiTokenPayload(BaseModel):
    account_id: str = Field(..., description="用户ID")
    tenant_id: str = Field(..., description="租户ID")
    # 限定只能是 dataset / app
    type: Literal["dataset", "app"] = Field(
        ..., description="令牌类型，仅支持 dataset、app"
    )
    # 非必填，默认 None
    app_id: Optional[str] = Field(None, description="应用ID，type 为 app 时必填")

    @field_validator("app_id")
    def check_app_id_required(cls, v, values):
        # values.data 获取模型已解析的其他字段
        token_type = values.data.get("type")
        if token_type == "app" and not v:
            raise ValueError("当令牌类型为 app 时，app_id 不能为空")
        if token_type == "dataset" and v is not None:
            raise ValueError("当令牌类型为 dataset 时，无需传入 app_id")
        return v

register_schema_models(
    console_ns,
    RegisterPayload,
    UserInfoPayload,
    GenerateApiTokenPayload,
)

register_response_schema_models(console_ns, SimpleResultOptionalDataResponse)

def get_first_tenant():
    stmt = select(Tenant).order_by(Tenant.created_at)
    return db.session.scalar(stmt)

# 对齐 Dify 原生密码哈希逻辑
def create_password_hash(raw_pwd: str) -> tuple[str, str]:
    salt = secrets.token_bytes(16)
    hashed = hash_password(raw_pwd, salt)
    return base64.b64encode(hashed).decode("utf-8"), base64.b64encode(salt).decode(
        "utf-8"
    )

@console_ns.route("/register")
class RegisterApi(Resource):
    @setup_required
    @private_api_auth_required
    @console_ns.expect(console_ns.models[RegisterPayload.__name__])
    @console_ns.response(
        200, "Success", console_ns.models[SimpleResultOptionalDataResponse.__name__]
    )
    @decrypt_password_field
    def post(self):
        # 单独捕获参数校验异常
        try:
            args = RegisterPayload.model_validate(console_ns.payload)
        except ValidationError as ve:
            err_info = ve.errors()[0]
            field = ".".join(map(str, err_info["loc"]))
            msg = f"Parameter error: {field} - {err_info['msg']}"
            return {"result": "fail", "message": msg}, 403

        try:
            email = args.email.strip().lower()
            name = args.name.strip()
            raw_pwd = args.password.strip()
            lang = get_valid_language(args.language)
            tz = args.timezone.strip()

            # 邮箱去重
            exist_stmt = select(Account).where(Account.email == email)
            if db.session.scalar(exist_stmt):
                raise AccountRegisterError("Email already exists")

            # 生成密码
            pwd_hash, pwd_salt = create_password_hash(raw_pwd)

            # 新建账号（字段完全对齐数据库）
            new_account = Account(
                email=email,
                name=name,
                password=pwd_hash,
                password_salt=pwd_salt,
                interface_language=lang,
                timezone=tz,
            )
            db.session.add(new_account)
            db.session.flush()

            # 租户处理
            tenant = get_first_tenant()
            if not tenant:
                tenant = Tenant(name=f"{name}'s Workspace")
                db.session.add(tenant)
                db.session.flush()

            is_first_tenant = get_first_tenant() is None
            join_rel = TenantAccountJoin(
                account_id=new_account.id,
                tenant_id=tenant.id,
                role="owner" if is_first_tenant else "admin",
                current=True,
            )
            db.session.add(join_rel)

            # 更新当前租户标记
            update_all = (
                update(TenantAccountJoin)
                .where(TenantAccountJoin.account_id == new_account.id)
                .values(current=False)
            )
            db.session.execute(update_all)

            update_curr = (
                update(TenantAccountJoin)
                .where(
                    TenantAccountJoin.account_id == new_account.id,
                    TenantAccountJoin.tenant_id == tenant.id,
                )
                .values(current=True)
            )
            db.session.execute(update_curr)

            # 统一提交，事务原子化
            db.session.commit()

            return make_response(
                {
                    "result": "success",
                    "data": {
                        "account_id": str(new_account.id),
                        "email": new_account.email,
                        "name": new_account.name,
                        "tenant_id": str(tenant.id),
                        "tenant_name": tenant.name,
                    },
                }
            )

        except Exception as e:
            db.session.rollback()
            logger.exception("Register failed, rollback transaction")
            if isinstance(e, AccountRegisterError):
                return {"result": "fail", "message": str(e)}, 400
            return {"result": "fail", "message": "Register failed"}, 500

# ===================== 获取用户基本信息 =====================
@console_ns.route("/user/info")
class UserInfoApi(Resource):
    @setup_required
    @private_api_auth_required
    def get(self):
        try:
            # 从 URL 查询参数取值，而非 payload
            account_id = request.args.get("account_id", "").strip()
            if not account_id:
                return {"result": "fail", "message": "account_id is required"}, 403
            args = UserInfoPayload(account_id=account_id)
        except ValidationError as ve:
            err_info = ve.errors()[0]
            field = ".".join(map(str, err_info["loc"]))
            msg = f"Parameter error: {field} - {err_info['msg']}"
            return {"result": "fail", "message": msg}, 403

        try:
            account_id = args.account_id.strip()
            # 1. 查询单条用户基础信息
            user_stmt = select(
                Account.id,
                Account.email,
                Account.name,
                Account.status,
                Account.timezone,
                Account.interface_language,
                Account.created_at,
            ).where(Account.id == account_id)

            user_row = db.session.execute(user_stmt).first()
            if not user_row:
                return {"result": "success", "data": None, "message": "用户不存在"}

            # 2. 查询该用户关联的租户信息
            tenant_stmt = (
                select(
                    TenantAccountJoin.account_id,
                    TenantAccountJoin.tenant_id,
                    TenantAccountJoin.role,
                    TenantAccountJoin.current,
                    Tenant.name.label("tenant_name"),
                    Tenant.created_at.label("tenant_created_at"),
                )
                .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
                .where(TenantAccountJoin.account_id == account_id)
            )
            tenant_row = db.session.execute(tenant_stmt).first()

            # 查询用户的 api_tokens
            api_token_stmt = select(
                ApiToken.id,
                ApiToken.tenant_id,
                ApiToken.account_id,
                ApiToken.type,
                ApiToken.token,
                ApiToken.app_id,
                ApiToken.last_used_at,
            ).where(ApiToken.account_id == account_id)
            api_token_rows = db.session.execute(api_token_stmt).all()

            api_token_list = []
            for row in api_token_rows:
                api_token_list.append(
                    {
                        "id": str(row.id),
                        "tenant_id": str(row.tenant_id),
                        "account_id": str(row.account_id),
                        "type": row.type,
                        "token": row.token,
                        "app_id": row.app_id,
                        "last_used_at": (
                            row.last_used_at.isoformat() if row.last_used_at else None
                        ),
                    }
                )

            # 3. 组装最终返回数据
            user_data = {
                "user_id": str(user_row.id),
                "email": user_row.email,
                "name": user_row.name,
                "timezone": user_row.timezone,
                "interface_language": user_row.interface_language,
                "created_at": (
                    user_row.created_at.isoformat() if user_row.created_at else None
                ),
                "tenant": tenant_row,
                "api_token_list": api_token_list,
            }
            return {"result": "success", "data": user_data}

        except Exception as e:
            logger.exception("Get user list failed")
            return {"result": "fail", "message": "Query failed"}, 500

# ===================== 创建 api key =====================
@console_ns.route("/generate/api_key")
class GenerateApiTokenApi(Resource):
    @setup_required
    @private_api_auth_required
    @console_ns.expect(console_ns.models[GenerateApiTokenPayload.__name__])
    def post(self):
        try:
            args = GenerateApiTokenPayload.model_validate(console_ns.payload)
        except ValidationError as ve:
            err_info = ve.errors()[0]
            field = ".".join(map(str, err_info["loc"]))
            msg = f"Parameter error: {field} - {err_info['msg']}"
            return {"result": "fail", "message": msg}, 403

        try:
            account_id = args.account_id.strip()
            tenant_id = args.tenant_id.strip()
            type = args.type
            app_id = args.app_id

            # 一个租户最多可以创建 10 个 api tokens(todo 要破掉)
            current_key_count = db.session.scalar(
                select(func.count(ApiToken.id)).where(ApiToken.tenant_id == tenant_id)
            )

            if current_key_count >= 10:
                logger.warning(
                    f"User {account_id} has reached maximum API key limit (10)"
                )
                return {"result": "fail", "message": "API key limit reached"}, 403

            # 2. 生成API Key（和Dify原生生成规则完全一致）
            token_prefix = "sk-"  # Dify原生用户API Key前缀
            key = ApiToken.generate_api_key(token_prefix, 24)

            # 3. 创建并保存API Token
            api_token = ApiToken()
            api_token.tenant_id = tenant_id
            api_token.account_id = account_id
            api_token.type = type
            api_token.token = key

            if app_id:
                api_token.app_id = app_id
            db.session.add(api_token)
            db.session.commit()

            return {"result": "success", "data": {"token": key}}

        except Exception as e:
            logger.exception("Get user list failed")
            return {"result": "fail", "message": "Query failed"}, 500

# ===================== 用户列表接口（含租户 + API Token） =====================
@console_ns.route("/user/list")
class UserListApi(Resource):
    @setup_required
    @private_api_auth_required
    def get(self):
        try:
            # 解析分页参数
            page = request.args.get("page", 1, type=int)
            page_size = request.args.get("page_size", 20, type=int)
            # 限制最大条数，防护大数据查询
            page_size = min(page_size, 100)
            offset = (page - 1) * page_size

            # 1. 查询用户总数
            total_stmt = select(func.count(Account.id))
            total = db.session.scalar(total_stmt) or 0

            # 2. 分页查询用户基础信息
            user_stmt = (
                select(
                    Account.id,
                    Account.email,
                    Account.name,
                    Account.timezone,
                    Account.interface_language,
                    Account.created_at,
                )
                .order_by(Account.created_at.desc())
                .offset(offset)
                .limit(page_size)
            )
            user_rows = db.session.execute(user_stmt).all()
            user_id_list = [row.id for row in user_rows]

            # 无用户直接返回空列表
            if not user_id_list:
                return {
                    "result": "success",
                    "data": {
                        "total": total,
                        "page": page,
                        "page_size": page_size,
                        "list": [],
                    },
                }

            # 3. 批量查询 用户-租户 关联信息（修复变量名笔误）
            tenant_stmt = (
                select(
                    TenantAccountJoin.account_id,
                    TenantAccountJoin.tenant_id,
                    TenantAccountJoin.role,
                    TenantAccountJoin.current,
                    Tenant.name.label("tenant_name"),
                    Tenant.created_at.label("tenant_created_at"),
                )
                .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
                .where(TenantAccountJoin.account_id.in_(user_id_list))
            )
            # 此处原代码写错：tenant_rows → tenant_stmt
            tenant_rows = db.session.execute(tenant_stmt).all()

            # 按用户ID分组租户数据
            tenant_map = {}
            for row in tenant_rows:
                aid = row.account_id
                if aid not in tenant_map:
                    tenant_map[aid] = []
                tenant_map[aid].append(
                    {
                        "tenant_id": str(row.tenant_id),
                        "tenant_name": row.tenant_name,
                        "role": row.role,
                        "is_current": row.current,
                        "tenant_created_at": (
                            row.tenant_created_at.isoformat()
                            if row.tenant_created_at
                            else None
                        ),
                    }
                )

            # 5. 组装最终返回数据
            user_list = []
            for user in user_rows:
                user_id = user.id
                user_list.append(
                    {
                        "user_id": str(user_id),
                        "email": user.email,
                        "name": user.name,
                        "timezone": user.timezone,
                        "interface_language": user.interface_language,
                        "created_at": (
                            user.created_at.isoformat() if user.created_at else None
                        ),
                        "tenant_list": tenant_map.get(user_id, []),
                        # "api_token_list": token_map.get(user_id, []),
                    }
                )

            return {
                "result": "success",
                "data": {
                    "total": total,
                    "page": page,
                    "page_size": page_size,
                    "list": user_list,
                },
            }

        except Exception as e:
            logger.exception("Get user list failed")
            return {"result": "fail", "message": "Query failed"}, 500

并且将该文件注册进去。修改 api/controllers/console/init.py, 添加 "register" 就好。

然后看看最近这几个功能修改涉及到的文件。

最后，如果有需要可以去打一个镜像。

docker compose build api