feat:新增oauth2.0登录

This commit is contained in:
npc0-hue
2025-04-17 14:34:26 +08:00
parent 385b3d5aaa
commit 64c7f3de25
26 changed files with 817 additions and 81 deletions
+2 -9
View File
@@ -45,16 +45,9 @@ def get_oauth_providers():
redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/google",
)
if not dify_config.OAUTH2_CLIENT_ID or not dify_config.OAUTH2_CLIENT_SECRET:
oa_oauth = None
else:
oa_oauth = OaOAuth(
client_id=dify_config.OAUTH2_CLIENT_ID,
client_secret=dify_config.OAUTH2_CLIENT_SECRET,
redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/oauth2",
)
oauth2 = OaOAuth(client_id='', client_secret='', redirect_uri='') # Extend: oauth2
OAUTH_PROVIDERS = {"github": github_oauth, "google": google_oauth, "oauth2": oa_oauth}
OAUTH_PROVIDERS = {"github": github_oauth, "google": google_oauth, "oauth2": oauth2}
return OAUTH_PROVIDERS
+87 -41
View File
@@ -1,10 +1,13 @@
import json
import logging # 二开部分,针对oa登录报错问题,记录返回的code
import urllib.parse
from dataclasses import dataclass
from typing import Optional
import requests
from flask import current_app
from configs import dify_config # Extend OAuto third-party login
from extensions.ext_database import db # Extend OAuto third-party login
from models.system_extend import SystemIntegrationExtend, SystemIntegrationClassify # Extend OAuto third-party login
@dataclass
@@ -135,35 +138,85 @@ class GoogleOAuth(OAuth):
return OAuthUserInfo(id=str(raw_info["sub"]), name="", email=raw_info["email"])
# Extend Start: OAuth2
class OaOAuth(OAuth):
_AUTH_URL = ""
_Host = ""
_TOKEN_URL = ""
_USER_INFO_URL = ""
def get_auto2_conf(self):
# oauth start
integration = db.session.query(SystemIntegrationExtend).filter(
SystemIntegrationExtend.classify == SystemIntegrationClassify.SYSTEM_INTEGRATION_OAUTH_TWO).first()
if integration is None or (integration and not integration.status):
return {
"integration": integration,
"config": {},
"passwd": ""
}
return {
"integration": integration,
"passwd": integration.decodeSecret(),
"config": json.loads(integration.config)
}
def extract_data(self, dictionary, path):
"""
从字典中提取指定路径的数据
支持通配符'*'获取列表中所有元素的特定字段
Args:
dictionary (dict): 源字典
path (str): 以点分隔的路径,如 "data.info.name""data.items.*.name"
Returns:
提取的数据
"""
parts = path.split('.')
current = dictionary
for i, part in enumerate(parts):
if part == '*' and isinstance(current, list):
# 处理列表中的每个元素
remainder = '.'.join(parts[i + 1:])
if remainder:
return [self.extract_data(item, remainder) for item in current]
else:
return current
elif isinstance(current, dict) and part in current:
current = current[part]
else:
return None
return current
def get_authorization_url(self, invite_token: Optional[str] = None):
params = {
"client_id": self.client_id,
"redirect_uri": self.redirect_uri,
}
with current_app.app_context():
self._Host = current_app.config.get("OAUTH2_CLIENT_URL")
self._TOKEN_URL = current_app.config.get("OAUTH2_TOKEN_URL")
self._USER_INFO_URL = current_app.config.get("OAUTH2_USER_INFO_URL")
return f"{self._Host}{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
auto2_conf = self.get_auto2_conf()
integration = auto2_conf.get('integration')
if integration is None:
return
# 构建查询参数
query_string = urllib.parse.urlencode({
'redirect_uri': dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/oauth2",
'client_id': integration.app_id,
})
# 构建完整URL
config = auto2_conf.get('config')
return f"{config.get('server_url')}{config.get('authorize_url')}?{query_string}"
def get_access_token(self, code: str):
auto2_conf = self.get_auto2_conf()
integration = auto2_conf.get('integration')
if integration is None:
return ""
data = {
"client_id": self.client_id,
"client_secret": self.client_secret,
"code": code,
"client_id": integration.app_id,
"grant_type": "authorization_code",
"redirect_uri": self.redirect_uri,
"client_secret": auto2_conf.get('passwd'),
"redirect_uri": dify_config.CONSOLE_API_URL + "/console/api/oauth/authorize/oauth2",
}
with current_app.app_context():
self._Host = current_app.config.get("OAUTH2_CLIENT_URL")
headers = {"Accept": "application/json"}
response = requests.post(self._Host + self._TOKEN_URL, data=data, headers=headers)
config = auto2_conf.get('config')
response = requests.post(f"{config.get('server_url')}{config.get('token_url')}", data=data, headers=headers)
response.encoding = "utf-8"
if response.status_code != 200:
return ""
@@ -176,42 +229,35 @@ class OaOAuth(OAuth):
return access_token
def get_raw_user_info(self, token: str):
with current_app.app_context():
self._Host = current_app.config.get("OAUTH2_CLIENT_URL")
auto2_conf = self.get_auto2_conf()
if auto2_conf.get('integration') is None:
return ""
config = auto2_conf.get('config')
headers = {"Authorization": f"Bearer {token}"}
response = requests.get(self._Host + self._USER_INFO_URL, headers=headers)
response = requests.get(f"{config.get('server_url')}{config.get('userinfo_url')}", headers=headers)
response.raise_for_status()
return response.json()
def _transform_user_info(self, raw_info: dict) -> OAuthUserInfo:
# 检查 raw_info 是否为空或为 None
if not raw_info or not isinstance(raw_info, dict):
auto2_conf = self.get_auto2_conf()
if not raw_info or not isinstance(raw_info, dict) or auto2_conf.get('integration') is None:
return OAuthUserInfo(
id="",
name="",
email="",
)
# 如果data为空说明报错了
data = raw_info.get("data")
if data == {} or data is None or not isinstance(data, dict):
code = raw_info.get("code", "")
msg = raw_info.get("info", "")
logging.info(f"raw_info {raw_info}")
return OAuthUserInfo(
id="",
name="",
email="",
)
username = data.get("username")
name = data.get("name")
email = data.get("email")
# 提取参数
config = auto2_conf.get('config')
name = self.extract_data(raw_info, config.get('user_name_field'))
email = self.extract_data(raw_info, config.get('user_email_field'))
username = self.extract_data(raw_info, config.get('user_id_field'))
if not username:
raise ValueError("OA系统返回用户数据格式不正确。请返回进行重新登录。")
raise ValueError("OAuth2返回用户数据格式不正确。请返回进行重新登录。")
return OAuthUserInfo(
id=str(username) if username is not None else None,
name=str(name) if name is not None else None,
email=str(email) if email is not None else None,
)
# Extend Stop: OAuth
+3
View File
@@ -8,6 +8,7 @@ class SystemIntegrationClassify:
SYSTEM_INTEGRATION_DINGTALK = 1 # 钉钉
SYSTEM_INTEGRATION_WEIXIN = 2 # 微信
SYSTEM_INTEGRATION_FEI_SU = 3 # 飞书
SYSTEM_INTEGRATION_OAUTH_TWO = 4 # OAuth2
class SystemIntegrationExtend(db.Model):
@@ -21,8 +22,10 @@ class SystemIntegrationExtend(db.Model):
status = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
corp_id = db.Column(db.String(120), nullable=True)
agent_id = db.Column(db.String(120), nullable=True)
app_id = db.Column(db.String(120), nullable=True)
app_key = db.Column(db.String(120), nullable=True)
app_secret = db.Column(db.Text, nullable=True)
config = db.Column(db.Text, nullable=True)
def decodeSecret(self):
if len(self.app_secret) == 0:
+10 -3
View File
@@ -1,3 +1,4 @@
import json
from enum import StrEnum
from pydantic import BaseModel, ConfigDict
@@ -83,6 +84,7 @@ class SystemFeatureModel(BaseModel):
is_email_setup: bool = False
license: LicenseModel = LicenseModel()
is_custom_auth2: str = "" # extend: Customizing AUTH2
is_custom_auth2_logout: str = "" # extend: Customizing AUTH2
ding_talk_client_id: str = "" # extend: DingTalk third-party login
ding_talk_corp_id: str = "" # extend: DingTalk sidebar login
ding_talk: bool = "" # extend: DingTalk sidebar login
@@ -134,15 +136,20 @@ class FeatureService:
system_features.is_allow_register = dify_config.ALLOW_REGISTER
system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE
system_features.is_email_setup = dify_config.MAIL_TYPE is not None and dify_config.MAIL_TYPE != ""
# extend start: Customizing AUTH2
system_features.is_custom_auth2 = dify_config.OAUTH2_CLIENT_URL
# extend stop: Customizing AUTH2
# extend start: DingTalk third-party login
for i in db.session.query(SystemIntegrationExtend).filter(SystemIntegrationExtend.status == True).all():
if i.classify == SystemIntegrationClassify.SYSTEM_INTEGRATION_DINGTALK:
system_features.ding_talk_client_id = i.app_key
system_features.ding_talk_corp_id = i.corp_id
system_features.ding_talk = i.status
# Extend: OAuth2 Start
elif i.classify == SystemIntegrationClassify.SYSTEM_INTEGRATION_OAUTH_TWO:
config = json.loads(i.config)
system_features.is_custom_auth2 = i.status
if "logout_url" in config.keys():
system_features.is_custom_auth2_logout = "{}{}".format(
config['server_url'], config['logout_url'])
# Extend: OAuth2 Stop
# extend stop: DingTalk third-party login
@classmethod