From 1819083a25dae0a4f080b10ae71cecd69cacbbaa Mon Sep 17 00:00:00 2001
From: "Jiangjie.Bai" <bugatti_it@163.com>
Date: Thu, 25 Aug 2022 15:04:45 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=20custom=20=E8=AE=A4?=
 =?UTF-8?q?=E8=AF=81=E6=A8=A1=E5=9D=97=E5=8A=A0=E8=BD=BD=E9=80=BB=E8=BE=91?=
 =?UTF-8?q?=EF=BC=8C=E5=88=A4=E6=96=ADMD5=E5=80=BC=EF=BC=8C=E5=90=AF?=
 =?UTF-8?q?=E5=8A=A8=E6=97=B6=E5=8F=AA=E5=8A=A0=E8=BD=BD=E4=B8=80=E6=AC=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/authentication/backends/custom.py | 31 ++++++++++++--------------
 apps/jumpserver/conf.py                |  1 +
 apps/jumpserver/settings/auth.py       | 20 ++++++++++++++++-
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/apps/authentication/backends/custom.py b/apps/authentication/backends/custom.py
index d0b309b83..c98d1742c 100644
--- a/apps/authentication/backends/custom.py
+++ b/apps/authentication/backends/custom.py
@@ -4,29 +4,25 @@ from common.utils import get_logger
 from django.contrib.auth import get_user_model
 from authentication.signals import user_auth_failed, user_auth_success
 
-
 from .base import JMSModelBackend
 
 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):
-    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):
-        if not settings.AUTH_CUSTOM:
-            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
+        return settings.AUTH_CUSTOM and callable(custom_authenticate_method)
 
     @staticmethod
     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):
         try:
-            authenticate = self.load_authenticate_method()
-            userinfo: dict = authenticate(username=username, password=password, **kwargs)
+            userinfo: dict = custom_authenticate_method(
+                username=username, password=password, **kwargs
+            )
             user, created = self.get_or_create_user_from_userinfo(userinfo)
         except Exception as e:
             logger.error('Custom authenticate error: {}'.format(e))
diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py
index c0aa257ee..5ec858821 100644
--- a/apps/jumpserver/conf.py
+++ b/apps/jumpserver/conf.py
@@ -225,6 +225,7 @@ class Config(dict):
 
         # Custom Config
         'AUTH_CUSTOM': False,
+        'AUTH_CUSTOM_FILE_MD5': '',
 
         # Auth LDAP settings
         'AUTH_LDAP': False,
diff --git a/apps/jumpserver/settings/auth.py b/apps/jumpserver/settings/auth.py
index 099b91de7..4aafdd35b 100644
--- a/apps/jumpserver/settings/auth.py
+++ b/apps/jumpserver/settings/auth.py
@@ -210,8 +210,26 @@ AUTHENTICATION_BACKENDS = [
     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
-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)