Underground-Digital/Workflow-Engine

1import logging2from datetime import datetime, timezone3from typing import Optional4 5import requests6from flask import current_app, redirect, request7from flask_restful import Resource8from werkzeug.exceptions import Unauthorized9 10from configs import dify_config11from constants.languages import languages12from events.tenant_event import tenant_was_created13from extensions.ext_database import db14from libs.helper import extract_remote_ip15from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo16from models import Account17from models.account import AccountStatus18from services.account_service import AccountService, RegisterService, TenantService19from services.errors.account import AccountNotFoundError20from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError21from services.feature_service import FeatureService22 23from .. import api24 25 26def get_oauth_providers():27    with current_app.app_context():28        if not dify_config.GITHUB_CLIENT_ID or not dify_config.GITHUB_CLIENT_SECRET:29            github_oauth = None30        else:31            github_oauth = GitHubOAuth(32                client_id=dify_config.GITHUB_CLIENT_ID,33                client_secret=dify_config.GITHUB_CLIENT_SECRET,34                redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/github",35            )36        if not dify_config.GOOGLE_CLIENT_ID or not dify_config.GOOGLE_CLIENT_SECRET:37            google_oauth = None38        else:39            google_oauth = GoogleOAuth(40                client_id=dify_config.GOOGLE_CLIENT_ID,41                client_secret=dify_config.GOOGLE_CLIENT_SECRET,42                redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/google",43            )44 45        OAUTH_PROVIDERS = {"github": github_oauth, "google": google_oauth}46        return OAUTH_PROVIDERS47 48 49class OAuthLogin(Resource):50    def get(self, provider: str):51        invite_token = request.args.get("invite_token") or None52        OAUTH_PROVIDERS = get_oauth_providers()53        with current_app.app_context():54            oauth_provider = OAUTH_PROVIDERS.get(provider)55            print(vars(oauth_provider))56        if not oauth_provider:57            return {"error": "Invalid provider"}, 40058 59        auth_url = oauth_provider.get_authorization_url(invite_token=invite_token)60        return redirect(auth_url)61 62 63class OAuthCallback(Resource):64    def get(self, provider: str):65        OAUTH_PROVIDERS = get_oauth_providers()66        with current_app.app_context():67            oauth_provider = OAUTH_PROVIDERS.get(provider)68        if not oauth_provider:69            return {"error": "Invalid provider"}, 40070 71        code = request.args.get("code")72        state = request.args.get("state")73        invite_token = None74        if state:75            invite_token = state76 77        try:78            token = oauth_provider.get_access_token(code)79            user_info = oauth_provider.get_user_info(token)80        except requests.exceptions.HTTPError as e:81            logging.exception(f"An error occurred during the OAuth process with {provider}: {e.response.text}")82            return {"error": "OAuth process failed"}, 40083 84        if invite_token and RegisterService.is_valid_invite_token(invite_token):85            invitation = RegisterService._get_invitation_by_token(token=invite_token)86            if invitation:87                invitation_email = invitation.get("email", None)88                if invitation_email != user_info.email:89                    return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Invalid invitation token.")90 91            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin/invite-settings?invite_token={invite_token}")92 93        try:94            account = _generate_account(provider, user_info)95        except AccountNotFoundError:96            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Account not found.")97        except (WorkSpaceNotFoundError, WorkSpaceNotAllowedCreateError):98            return redirect(99                f"{dify_config.CONSOLE_WEB_URL}/signin"100                "?message=Workspace not found, please contact system admin to invite you to join in a workspace."101            )102 103        # Check account status104        if account.status == AccountStatus.BANNED.value:105            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Account is banned.")106 107        if account.status == AccountStatus.PENDING.value:108            account.status = AccountStatus.ACTIVE.value109            account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None)110            db.session.commit()111 112        try:113            TenantService.create_owner_tenant_if_not_exist(account)114        except Unauthorized:115            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found.")116        except WorkSpaceNotAllowedCreateError:117            return redirect(118                f"{dify_config.CONSOLE_WEB_URL}/signin"119                "?message=Workspace not found, please contact system admin to invite you to join in a workspace."120            )121 122        token_pair = AccountService.login(123            account=account,124            ip_address=extract_remote_ip(request),125        )126 127        return redirect(128            f"{dify_config.CONSOLE_WEB_URL}?access_token={token_pair.access_token}&refresh_token={token_pair.refresh_token}"129        )130 131 132def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) -> Optional[Account]:133    account = Account.get_by_openid(provider, user_info.id)134 135    if not account:136        account = Account.query.filter_by(email=user_info.email).first()137 138    return account139 140 141def _generate_account(provider: str, user_info: OAuthUserInfo):142    # Get account by openid or email.143    account = _get_account_by_openid_or_email(provider, user_info)144 145    if account:146        tenant = TenantService.get_join_tenants(account)147        if not tenant:148            if not FeatureService.get_system_features().is_allow_create_workspace:149                raise WorkSpaceNotAllowedCreateError()150            else:151                tenant = TenantService.create_tenant(f"{account.name}'s Workspace")152                TenantService.create_tenant_member(tenant, account, role="owner")153                account.current_tenant = tenant154                tenant_was_created.send(tenant)155 156    if not account:157        if not FeatureService.get_system_features().is_allow_register:158            raise AccountNotFoundError()159        account_name = user_info.name or "Dify"160        account = RegisterService.register(161            email=user_info.email, name=account_name, password=None, open_id=user_info.id, provider=provider162        )163 164        # Set interface language165        preferred_lang = request.accept_languages.best_match(languages)166        if preferred_lang and preferred_lang in languages:167            interface_language = preferred_lang168        else:169            interface_language = languages[0]170        account.interface_language = interface_language171        db.session.commit()172 173    # Link account174    AccountService.link_account_integrate(provider, user_info.id, account)175 176    return account177 178 179api.add_resource(OAuthLogin, "/oauth/login/<provider>")180api.add_resource(OAuthCallback, "/oauth/authorize/<provider>")181