feat: add oauth invite redict

This commit is contained in:
Joe 2024-09-02 14:50:45 +08:00
parent cbdbfb844d
commit c84f004035
2 changed files with 19 additions and 4 deletions

View File

@ -43,7 +43,7 @@ def get_oauth_providers():
class OAuthLogin(Resource):
def get(self, provider: str):
def get(self, provider: str, invite_toke: Optional[str] = None):
OAUTH_PROVIDERS = get_oauth_providers()
with current_app.app_context():
oauth_provider = OAUTH_PROVIDERS.get(provider)
@ -51,7 +51,7 @@ class OAuthLogin(Resource):
if not oauth_provider:
return {"error": "Invalid provider"}, 400
auth_url = oauth_provider.get_authorization_url()
auth_url = oauth_provider.get_authorization_url(invite_toke)
return redirect(auth_url)
@ -64,6 +64,11 @@ class OAuthCallback(Resource):
return {"error": "Invalid provider"}, 400
code = request.args.get("code")
state = request.args.get("state")
invite_token = None
if state:
invite_token = state
try:
token = oauth_provider.get_access_token(code)
user_info = oauth_provider.get_user_info(token)
@ -71,6 +76,11 @@ class OAuthCallback(Resource):
logging.exception(f"An error occurred during the OAuth process with {provider}: {e.response.text}")
return {"error": "OAuth process failed"}, 400
if invite_token:
return redirect(
f"{dify_config.CONSOLE_WEB_URL}/invite-settings?invite_token={invite_token}"
)
try:
account = _generate_account(provider, user_info)
except services.errors.account.AccountNotFound as e:

View File

@ -1,5 +1,6 @@
import urllib.parse
from dataclasses import dataclass
from typing import Optional
import requests
@ -40,12 +41,14 @@ class GitHubOAuth(OAuth):
_USER_INFO_URL = "https://api.github.com/user"
_EMAIL_INFO_URL = "https://api.github.com/user/emails"
def get_authorization_url(self):
def get_authorization_url(self, invite_token: Optional[str] = None):
params = {
"client_id": self.client_id,
"redirect_uri": self.redirect_uri,
"scope": "user:email", # Request only basic user information
}
if invite_token:
params["state"] = invite_token
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
def get_access_token(self, code: str):
@ -90,13 +93,15 @@ class GoogleOAuth(OAuth):
_TOKEN_URL = "https://oauth2.googleapis.com/token"
_USER_INFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo"
def get_authorization_url(self):
def get_authorization_url(self, invite_token: Optional[str] = None):
params = {
"client_id": self.client_id,
"response_type": "code",
"redirect_uri": self.redirect_uri,
"scope": "openid email",
}
if invite_token:
params["state"] = invite_token
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
def get_access_token(self, code: str):