dify/api/controllers/cli_api/dify_cli/wraps.py

from collections.abc import Callable
from functools import wraps
from typing import ParamSpec, TypeVar

from flask import current_app, g, request
from flask_login import user_logged_in
from pydantic import BaseModel
from sqlalchemy.orm import Session

from core.session.cli_api import CliApiSession, CliContext
from extensions.ext_database import db
from libs.login import current_user
from models.account import Tenant
from models.model import DefaultEndUserSessionID, EndUser

P = ParamSpec("P")
R = TypeVar("R")


class TenantUserPayload(BaseModel):
    tenant_id: str
    user_id: str


def get_user(tenant_id: str, user_id: str | None) -> EndUser:
    """
    Get current user

    NOTE: user_id is not trusted, it could be maliciously set to any value.
    As a result, it could only be considered as an end user id.
    """
    if not user_id:
        user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
    is_anonymous = user_id == DefaultEndUserSessionID.DEFAULT_SESSION_ID
    try:
        with Session(db.engine) as session:
            user_model = None

            if is_anonymous:
                user_model = (
                    session.query(EndUser)
                    .where(
                        EndUser.session_id == user_id,
                        EndUser.tenant_id == tenant_id,
                    )
                    .first()
                )
            else:
                user_model = (
                    session.query(EndUser)
                    .where(
                        EndUser.id == user_id,
                        EndUser.tenant_id == tenant_id,
                    )
                    .first()
                )

            if not user_model:
                user_model = EndUser(
                    tenant_id=tenant_id,
                    type="service_api",
                    is_anonymous=is_anonymous,
                    session_id=user_id,
                )
                session.add(user_model)
                session.commit()
                session.refresh(user_model)

    except Exception:
        raise ValueError("user not found")

    return user_model


def get_cli_user_tenant(view_func: Callable[P, R]):
    @wraps(view_func)
    def decorated_view(*args: P.args, **kwargs: P.kwargs):
        session: CliApiSession | None = getattr(g, "cli_api_session", None)
        if session is None:
            raise ValueError("session not found")

        user_id = session.user_id
        tenant_id = session.tenant_id
        cli_context = CliContext.model_validate(session.context)

        if not user_id:
            user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID

        try:
            tenant_model = (
                db.session.query(Tenant)
                .where(
                    Tenant.id == tenant_id,
                )
                .first()
            )
        except Exception:
            raise ValueError("tenant not found")

        if not tenant_model:
            raise ValueError("tenant not found")

        kwargs["tenant_model"] = tenant_model
        kwargs["user_model"] = get_user(tenant_id, user_id)
        kwargs["cli_context"] = cli_context

        current_app.login_manager._update_request_context_with_user(kwargs["user_model"])  # type: ignore
        user_logged_in.send(current_app._get_current_object(), user=current_user)  # type: ignore

        return view_func(*args, **kwargs)

    return decorated_view


def plugin_data(view: Callable[P, R] | None = None, *, payload_type: type[BaseModel]):
    def decorator(view_func: Callable[P, R]):
        @wraps(view_func)
        def decorated_view(*args: P.args, **kwargs: P.kwargs):
            try:
                data = request.get_json()
            except Exception:
                raise ValueError("invalid json")

            try:
                payload = payload_type.model_validate(data)
            except Exception as e:
                raise ValueError(f"invalid payload: {str(e)}")

            kwargs["payload"] = payload
            return view_func(*args, **kwargs)

        return decorated_view

    if view is None:
        return decorator
    else:
        return decorator(view)