Merge pull request #8811 from jumpserver/pr@dev@perf_customauth

perf: 优化 custom 认证模块加载逻辑,判断MD5值,启动时只加载一次
pull/8816/head
老广 2022-08-25 15:23:08 +08:00 committed by GitHub
commit 2cb6da3129
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 34 additions and 18 deletions

View File

@ -4,29 +4,25 @@ from common.utils import get_logger
from django.contrib.auth import get_user_model from django.contrib.auth import get_user_model
from authentication.signals import user_auth_failed, user_auth_success from authentication.signals import user_auth_failed, user_auth_success
from .base import JMSModelBackend from .base import JMSModelBackend
logger = get_logger(__file__) logger = get_logger(__file__)
custom_authenticate_method = None
if settings.AUTH_CUSTOM:
""" 保证自定义认证方法在服务运行时不能被更改,只在第一次调用时加载一次 """
try:
custom_auth_method_path = 'data.auth.main.authenticate'
custom_authenticate_method = import_string(custom_auth_method_path)
except Exception as e:
logger.warning('Import custom auth method failed: {}, Maybe not enabled'.format(e))
class CustomAuthBackend(JMSModelBackend): class CustomAuthBackend(JMSModelBackend):
custom_auth_method_path = 'data.auth.main.authenticate'
def load_authenticate_method(self):
return import_string(self.custom_auth_method_path)
def is_enabled(self): def is_enabled(self):
if not settings.AUTH_CUSTOM: return settings.AUTH_CUSTOM and callable(custom_authenticate_method)
return False
try:
self.load_authenticate_method()
except Exception as e:
logger.warning('Not enabled custom auth backend: {}'.format(e))
return False
else:
logger.info('Enabled custom auth backend')
return True
@staticmethod @staticmethod
def get_or_create_user_from_userinfo(userinfo: dict): def get_or_create_user_from_userinfo(userinfo: dict):
@ -40,8 +36,9 @@ class CustomAuthBackend(JMSModelBackend):
def authenticate(self, request, username=None, password=None, **kwargs): def authenticate(self, request, username=None, password=None, **kwargs):
try: try:
authenticate = self.load_authenticate_method() userinfo: dict = custom_authenticate_method(
userinfo: dict = authenticate(username=username, password=password, **kwargs) username=username, password=password, **kwargs
)
user, created = self.get_or_create_user_from_userinfo(userinfo) user, created = self.get_or_create_user_from_userinfo(userinfo)
except Exception as e: except Exception as e:
logger.error('Custom authenticate error: {}'.format(e)) logger.error('Custom authenticate error: {}'.format(e))

View File

@ -225,6 +225,7 @@ class Config(dict):
# Custom Config # Custom Config
'AUTH_CUSTOM': False, 'AUTH_CUSTOM': False,
'AUTH_CUSTOM_FILE_MD5': '',
# Auth LDAP settings # Auth LDAP settings
'AUTH_LDAP': False, 'AUTH_LDAP': False,

View File

@ -210,8 +210,26 @@ AUTHENTICATION_BACKENDS = [
AUTH_BACKEND_AUTH_TOKEN, AUTH_BACKEND_SSO, AUTH_BACKEND_TEMP_TOKEN, AUTH_BACKEND_AUTH_TOKEN, AUTH_BACKEND_SSO, AUTH_BACKEND_TEMP_TOKEN,
] ]
def get_file_md5(filepath):
import hashlib
# 创建md5对象
m = hashlib.md5()
with open(filepath, 'rb') as f:
while True:
data = f.read(4096)
if not data:
break
# 更新md5对象
m.update(data)
# 返回md5对象
return m.hexdigest()
AUTH_CUSTOM = CONFIG.AUTH_CUSTOM AUTH_CUSTOM = CONFIG.AUTH_CUSTOM
if AUTH_CUSTOM: AUTH_CUSTOM_FILE_MD5 = CONFIG.AUTH_CUSTOM_FILE_MD5
AUTH_CUSTOM_FILE_PATH = os.path.join(PROJECT_DIR, 'data', 'auth', 'main.py')
if AUTH_CUSTOM and AUTH_CUSTOM_FILE_MD5 == get_file_md5(AUTH_CUSTOM_FILE_PATH):
# 自定义认证模块 # 自定义认证模块
AUTHENTICATION_BACKENDS.append(AUTH_BACKEND_CUSTOM) AUTHENTICATION_BACKENDS.append(AUTH_BACKEND_CUSTOM)