From e2d5b69510ce47cdf8a3ac569ecd4cea378fcf8b Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Tue, 13 Apr 2021 16:21:51 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E5=81=A5=E5=BA=B7?=
 =?UTF-8?q?=E7=9B=91=E6=B5=8B=EF=BC=8C=E5=B9=B6=E6=B7=BB=E5=8A=A0=20health?=
 =?UTF-8?q?=20check=20=E7=9A=84=20key?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/jumpserver/api.py             | 67 +++++++++++++++++++++++++++++-
 apps/jumpserver/conf.py            |  1 +
 apps/jumpserver/settings/custom.py |  1 +
 apps/jumpserver/urls.py            |  3 +-
 apps/jumpserver/views/other.py     |  9 +---
 5 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/apps/jumpserver/api.py b/apps/jumpserver/api.py
index e31a1d843..bda2537c8 100644
--- a/apps/jumpserver/api.py
+++ b/apps/jumpserver/api.py
@@ -1,3 +1,5 @@
+import time
+
 from django.core.cache import cache
 from django.utils import timezone
 from django.utils.timesince import timesince
@@ -6,6 +8,8 @@ from django.http.response import JsonResponse, HttpResponse
 from rest_framework.views import APIView
 from rest_framework.permissions import AllowAny
 from collections import Counter
+from django.conf import settings
+from rest_framework.response import Response
 
 from users.models import User
 from assets.models import Asset
@@ -307,7 +311,68 @@ class IndexApi(TotalCountMixin, DatesLoginMetricMixin, APIView):
         return JsonResponse(data, status=200)
 
 
-class PrometheusMetricsApi(APIView):
+class HealthApiMixin(APIView):
+    def is_token_right(self):
+        token = self.request.query_params.get('token')
+        ok_token = settings.HEALTH_CHECK_TOKEN
+        if ok_token and token != ok_token:
+            return False
+        return True
+
+    def check_permissions(self, request):
+        if not self.is_token_right():
+            msg = 'Health check token error, ' \
+                  'Please set query param in url and same with setting HEALTH_CHECK_TOKEN. ' \
+                  'eg: $PATH/?token=$HEALTH_CHECK_TOKEN'
+            self.permission_denied(request, message={'error': msg}, code=403)
+
+
+class HealthCheckView(HealthApiMixin):
+    permission_classes = (AllowAny,)
+
+    @staticmethod
+    def get_db_status():
+        t1 = time.time()
+        try:
+            User.objects.first()
+            t2 = time.time()
+            return True, t2 - t1
+        except:
+            t2 = time.time()
+            return False, t2 - t1
+
+    def get_redis_status(self):
+        key = 'HEALTH_CHECK'
+
+        t1 = time.time()
+        try:
+            value = '1'
+            cache.set(key, '1', 10)
+            got = cache.get(key)
+            t2 = time.time()
+            if value == got:
+                return True, t2 -t1
+            return False, t2 -t1
+        except:
+            t2 = time.time()
+            return False, t2 - t1
+
+    def get(self, request):
+        redis_status, redis_time = self.get_redis_status()
+        db_status, db_time = self.get_db_status()
+        status = all([redis_status, db_status])
+        data = {
+            'status': status,
+            'db_status': db_status,
+            'db_time': db_time,
+            'redis_status': redis_status,
+            'redis_time': redis_time,
+            'time': int(time.time())
+        }
+        return Response(data)
+
+
+class PrometheusMetricsApi(HealthApiMixin):
     permission_classes = (AllowAny,)
 
     def get(self, request, *args, **kwargs):
diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py
index d09a81bd2..179376828 100644
--- a/apps/jumpserver/conf.py
+++ b/apps/jumpserver/conf.py
@@ -289,6 +289,7 @@ class Config(dict):
         'SESSION_SAVE_EVERY_REQUEST': True,
         'SESSION_EXPIRE_AT_BROWSER_CLOSE_FORCE': False,
         'FORGOT_PASSWORD_URL': '',
+        'HEALTH_CHECK_TOKEN': ''
     }
 
     def compatible_auth_openid_of_key(self):
diff --git a/apps/jumpserver/settings/custom.py b/apps/jumpserver/settings/custom.py
index 89b8d6d53..936b27582 100644
--- a/apps/jumpserver/settings/custom.py
+++ b/apps/jumpserver/settings/custom.py
@@ -123,3 +123,4 @@ FORGOT_PASSWORD_URL = CONFIG.FORGOT_PASSWORD_URL
 
 # 自定义默认组织名
 GLOBAL_ORG_DISPLAY_NAME = CONFIG.GLOBAL_ORG_DISPLAY_NAME
+HEALTH_CHECK_TOKEN = CONFIG.HEALTH_CHECK_TOKEN
diff --git a/apps/jumpserver/urls.py b/apps/jumpserver/urls.py
index 044d09310..759c6f271 100644
--- a/apps/jumpserver/urls.py
+++ b/apps/jumpserver/urls.py
@@ -48,7 +48,8 @@ urlpatterns = [
     path('', views.IndexView.as_view(), name='index'),
     path('api/v1/', include(api_v1)),
     re_path('api/(?P<app>\w+)/(?P<version>v\d)/.*', views.redirect_format_api),
-    path('api/health/', views.HealthCheckView.as_view(), name="health"),
+    path('api/health/', api.HealthCheckView.as_view(), name="health"),
+    path('api/v1/health/', api.HealthCheckView.as_view(), name="health_v1"),
     # External apps url
     path('core/auth/captcha/', include('captcha.urls')),
     path('core/', include(app_view_patterns)),
diff --git a/apps/jumpserver/views/other.py b/apps/jumpserver/views/other.py
index da8046bfc..9ab561c4c 100644
--- a/apps/jumpserver/views/other.py
+++ b/apps/jumpserver/views/other.py
@@ -17,7 +17,7 @@ from common.http import HttpResponseTemporaryRedirect
 
 
 __all__ = [
-    'LunaView', 'I18NView', 'KokoView', 'WsView', 'HealthCheckView',
+    'LunaView', 'I18NView', 'KokoView', 'WsView',
     'redirect_format_api', 'redirect_old_apps_view', 'UIView'
 ]
 
@@ -64,13 +64,6 @@ def redirect_old_apps_view(request, *args, **kwargs):
     return HttpResponseTemporaryRedirect(new_path)
 
 
-class HealthCheckView(APIView):
-    permission_classes = (AllowAny,)
-
-    def get(self, request):
-        return JsonResponse({"status": 1, "time": int(time.time())})
-
-
 class WsView(APIView):
     ws_port = settings.HTTP_LISTEN_PORT + 1