From d8d487f770c65da67267b3c7128a9e347169d9ab Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Mon, 24 Jul 2023 15:32:30 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BF=AE=E6=94=B9=20ALLOW=5FHOSTS?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/utils/encode.py      |  2 +-
 apps/common/utils/timezone.py    | 15 +--------------
 apps/jumpserver/asgi.py          |  4 ++--
 apps/jumpserver/conf.py          |  1 +
 apps/jumpserver/settings/base.py | 18 +++++++++++++-----
 5 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/apps/common/utils/encode.py b/apps/common/utils/encode.py
index 6f1fc079e..767645c88 100644
--- a/apps/common/utils/encode.py
+++ b/apps/common/utils/encode.py
@@ -45,7 +45,7 @@ class Signer(metaclass=Singleton):
 
     def sign(self, value):
         s = JSONWebSignatureSerializer(self.secret_key, algorithm_name='HS256')
-        return self.json_serializer.dumps(value).decode()
+        return s.dumps(value).decode()
 
     def unsign(self, value):
         if value is None:
diff --git a/apps/common/utils/timezone.py b/apps/common/utils/timezone.py
index b7dfd595f..a74ebe73a 100644
--- a/apps/common/utils/timezone.py
+++ b/apps/common/utils/timezone.py
@@ -1,21 +1,8 @@
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 
-import pytz
 from django.utils import timezone as dj_timezone
 from rest_framework.fields import DateTimeField
 
-max = datetime.max.replace(tzinfo=timezone.utc)
-
-
-def astimezone(dt: datetime, tzinfo: pytz.tzinfo.DstTzInfo):
-    assert dj_timezone.is_aware(dt)
-    print("dt.tzinfo: ", tzinfo, type(tzinfo))
-    return tzinfo.normalize(dt.astimezone(tzinfo))
-
-
-def as_china_cst(dt: datetime):
-    return astimezone(dt, pytz.timezone('Asia/Shanghai'))
-
 
 def as_current_tz(dt: datetime):
     return dt.astimezone(dj_timezone.get_current_timezone())
diff --git a/apps/jumpserver/asgi.py b/apps/jumpserver/asgi.py
index 428781225..dc7760efe 100644
--- a/apps/jumpserver/asgi.py
+++ b/apps/jumpserver/asgi.py
@@ -2,9 +2,9 @@ import os
 
 from channels.auth import AuthMiddlewareStack
 from channels.routing import ProtocolTypeRouter, URLRouter
-from channels.security.websocket import AllowedHostsOriginValidator
 from django.core.asgi import get_asgi_application
 
+from .middleware import WsSignatureAuthMiddleware
 from .routing import urlpatterns
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
@@ -13,7 +13,7 @@ application = ProtocolTypeRouter({
     "http": get_asgi_application(),
 
     # WebSocket chat handler
-    "websocket": AllowedHostsOriginValidator(
+    "websocket": WsSignatureAuthMiddleware(
         AuthMiddlewareStack(
             URLRouter(urlpatterns)
         )
diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py
index d9b15943f..711edf657 100644
--- a/apps/jumpserver/conf.py
+++ b/apps/jumpserver/conf.py
@@ -514,6 +514,7 @@ class Config(dict):
         'TIME_ZONE': 'Asia/Shanghai',
         'FORCE_SCRIPT_NAME': '',
         'SESSION_COOKIE_SECURE': False,
+        'ALLOWED_HOSTS': '',
         'CSRF_COOKIE_SECURE': False,
         'REFERER_CHECK_ENABLED': False,
         'CSRF_TRUSTED_ORIGINS': '',
diff --git a/apps/jumpserver/settings/base.py b/apps/jumpserver/settings/base.py
index 179fe77af..551e4ae8a 100644
--- a/apps/jumpserver/settings/base.py
+++ b/apps/jumpserver/settings/base.py
@@ -65,14 +65,22 @@ APPLET_DOWNLOAD_HOST = CONFIG.APPLET_DOWNLOAD_HOST
 # https://docs.djangoproject.com/en/4.1/ref/settings/
 SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https')
 
-# https://docs.djangoproject.com/en/4.1/ref/settings/#std-setting-CSRF_TRUSTED_ORIGINS
-CSRF_TRUSTED_ORIGINS = CONFIG.CSRF_TRUSTED_ORIGINS.split(',') if CONFIG.CSRF_TRUSTED_ORIGINS \
-    else ['https://*', 'https://.*', 'http://localhost:9528']
-
 # LOG LEVEL
 LOG_LEVEL = CONFIG.LOG_LEVEL
 
-ALLOWED_HOSTS = ['*']
+ALLOWED_HOSTS = CONFIG.ALLOWED_HOSTS.split(',') if CONFIG.ALLOWED_HOSTS else ['localhost', '127.0.0.1']
+
+# https://docs.djangoproject.com/en/4.1/ref/settings/#std-setting-CSRF_TRUSTED_ORIGINS
+CSRF_TRUSTED_ORIGINS = []
+for origin in ALLOWED_HOSTS:
+    # 避免错误 先判断一下吧
+    if origin.startswith('http'):
+        CSRF_TRUSTED_ORIGINS.append(origin)
+        continue
+    if origin.startswith('.'):
+        origin = '*.'
+    for schema in ['https', 'http']:
+        CSRF_TRUSTED_ORIGINS.append('{}://{}'.format(schema, origin))
 
 # Max post update field num
 DATA_UPLOAD_MAX_NUMBER_FIELDS = 10000