dify/api/core/plugin/impl/oauth.py

import binascii
from collections.abc import Mapping
from typing import Any

from werkzeug import Request

from core.dify_builder.real_call_consent import validate_confirmed_dispatch_headers
from core.plugin.entities.plugin_daemon import PluginOAuthAuthorizationUrlResponse, PluginOAuthCredentialsResponse
from core.plugin.impl.base import BasePluginClient


class OAuthHandler(BasePluginClient):
    def get_authorization_url(
        self,
        tenant_id: str,
        user_id: str,
        plugin_id: str,
        provider: str,
        redirect_uri: str,
        system_credentials: Mapping[str, Any],
    ) -> PluginOAuthAuthorizationUrlResponse:
        try:
            response = self._request_with_plugin_daemon_response_stream(
                "POST",
                f"plugin/{tenant_id}/dispatch/oauth/get_authorization_url",
                PluginOAuthAuthorizationUrlResponse,
                data={
                    "user_id": user_id,
                    "data": {
                        "provider": provider,
                        "redirect_uri": redirect_uri,
                        "system_credentials": system_credentials,
                    },
                },
                headers={
                    "X-Plugin-ID": plugin_id,
                    "Content-Type": "application/json",
                },
            )
            for resp in response:
                return resp
            raise ValueError("No response received from plugin daemon for authorization URL request.")
        except Exception as e:
            raise ValueError(f"Error getting authorization URL: {e}")

    def get_credentials(
        self,
        tenant_id: str,
        user_id: str,
        plugin_id: str,
        provider: str,
        redirect_uri: str,
        system_credentials: Mapping[str, Any],
        request: Request,
    ) -> PluginOAuthCredentialsResponse:
        """
        Get credentials from the given request.
        """

        try:
            # encode request to raw http request
            raw_request_bytes = self._convert_request_to_raw_data(request)
            response = self._request_with_plugin_daemon_response_stream(
                "POST",
                f"plugin/{tenant_id}/dispatch/oauth/get_credentials",
                PluginOAuthCredentialsResponse,
                data={
                    "user_id": user_id,
                    "data": {
                        "provider": provider,
                        "redirect_uri": redirect_uri,
                        "system_credentials": system_credentials,
                        # for json serialization
                        "raw_http_request": binascii.hexlify(raw_request_bytes).decode(),
                    },
                },
                headers={
                    "X-Plugin-ID": plugin_id,
                    "Content-Type": "application/json",
                },
            )
            for resp in response:
                return resp
            raise ValueError("No response received from plugin daemon for authorization URL request.")
        except Exception as e:
            raise ValueError(f"Error getting credentials: {e}")

    def refresh_credentials(
        self,
        tenant_id: str,
        user_id: str,
        plugin_id: str,
        provider: str,
        redirect_uri: str,
        system_credentials: Mapping[str, Any],
        credentials: Mapping[str, Any],
        *,
        confirmed_dispatch_headers: Mapping[str, str] | None = None,
    ) -> PluginOAuthCredentialsResponse:
        """Refresh through the native tenant route.

        ``confirmed_dispatch_headers`` is ``None`` for every ordinary refresh. A
        confirmed Builder refresh supplies exactly the three server-built
        companion headers, validated before any transport.
        """
        headers = {
            "X-Plugin-ID": plugin_id,
            "Content-Type": "application/json",
        }
        if confirmed_dispatch_headers is not None:
            headers.update(validate_confirmed_dispatch_headers(confirmed_dispatch_headers))
        try:
            response = self._request_with_plugin_daemon_response_stream(
                "POST",
                f"plugin/{tenant_id}/dispatch/oauth/refresh_credentials",
                PluginOAuthCredentialsResponse,
                data={
                    "user_id": user_id,
                    "data": {
                        "provider": provider,
                        "redirect_uri": redirect_uri,
                        "system_credentials": system_credentials,
                        "credentials": credentials,
                    },
                },
                headers=headers,
            )
            for resp in response:
                return resp
            raise ValueError("No response received from plugin daemon for refresh credentials request.")
        except Exception as e:
            raise ValueError(f"Error refreshing credentials: {e}")

    def _convert_request_to_raw_data(self, request: Request) -> bytes:
        """
        Convert a Request object to raw HTTP data.

        Args:
            request: The Request object to convert.

        Returns:
            The raw HTTP data as bytes.
        """
        # Start with the request line
        method = request.method
        path = request.full_path
        protocol = request.headers.get("HTTP_VERSION", "HTTP/1.1")
        raw_data = f"{method} {path} {protocol}\r\n".encode()

        # Add headers
        for header_name, header_value in request.headers.items():
            raw_data += f"{header_name}: {header_value}\r\n".encode()

        # Add empty line to separate headers from body
        raw_data += b"\r\n"

        # Add body if exists
        body = request.get_data(as_text=False)
        if body:
            raw_data += body

        return raw_data