From 2cf80e6615771168be18ab09127a187921954959 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Thu, 31 Oct 2024 18:22:43 +0800
Subject: [PATCH] perf: login success to call client

---
 apps/authentication/views/login.py           |  2 +-
 apps/templates/flash_message_standalone.html |  5 ++--
 apps/users/models/user/_token.py             |  4 +--
 apps/users/utils.py                          | 28 +++++++++++++++++++-
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/apps/authentication/views/login.py b/apps/authentication/views/login.py
index f487f3db7..e43969ca3 100644
--- a/apps/authentication/views/login.py
+++ b/apps/authentication/views/login.py
@@ -24,8 +24,8 @@ from django.views.decorators.debug import sensitive_post_parameters
 from django.views.generic.base import TemplateView, RedirectView
 from django.views.generic.edit import FormView
 
-from common.utils import FlashMessageUtil, static_or_direct, safe_next_url
 from common.const import Language
+from common.utils import FlashMessageUtil, static_or_direct, safe_next_url
 from users.utils import (
     redirect_user_first_login_or_index
 )
diff --git a/apps/templates/flash_message_standalone.html b/apps/templates/flash_message_standalone.html
index fd3308f00..da34e1aae 100644
--- a/apps/templates/flash_message_standalone.html
+++ b/apps/templates/flash_message_standalone.html
@@ -16,7 +16,7 @@
             {% if error %}
                 {{ error }}
             {% else %}
-                {{ message|safe }}
+                {{ message | safe }}
             {% endif %}
             </div>
         </p>
@@ -52,6 +52,7 @@
     {% else %}
         message = '{{ message|safe }}'
     {% endif %}
+    var redirect_url = '{{ redirect_url }}'
 
     function redirect_page() {
         if (time >= 0) {
@@ -60,7 +61,7 @@
             time -= ttl;
             setTimeout(redirect_page, ttl * 1000);
         } else {
-            window.location.href = "{{ redirect_url }}";
+            window.location.href = redirect_url
         }
     }
     {% if auto_redirect %}
diff --git a/apps/users/models/user/_token.py b/apps/users/models/user/_token.py
index 8a568316c..2c3991e37 100644
--- a/apps/users/models/user/_token.py
+++ b/apps/users/models/user/_token.py
@@ -36,8 +36,8 @@ class TokenMixin:
         self.delete_private_token()
         return self.create_private_token()
 
-    def create_bearer_token(self, request=None):
-        expiration = settings.TOKEN_EXPIRATION or 3600
+    def create_bearer_token(self, request=None, age=None):
+        expiration = age or settings.TOKEN_EXPIRATION or 3600
         if request:
             remote_addr = request.META.get("REMOTE_ADDR", "")
         else:
diff --git a/apps/users/utils.py b/apps/users/utils.py
index 14b36a875..9893b331e 100644
--- a/apps/users/utils.py
+++ b/apps/users/utils.py
@@ -1,6 +1,7 @@
 # ~*~ coding: utf-8 ~*~
 #
 import base64
+import json
 import logging
 import os
 import re
@@ -9,9 +10,10 @@ import time
 import pyotp
 from django.conf import settings
 from django.core.cache import cache
+from django.utils.translation import gettext as _
 
 from common.tasks import send_mail_async
-from common.utils import reverse, get_object_or_none, ip, safe_next_url
+from common.utils import reverse, get_object_or_none, ip, safe_next_url, FlashMessageUtil
 from .models import User
 
 logger = logging.getLogger('jumpserver.users')
@@ -45,10 +47,34 @@ def get_user_or_pre_auth_user(request):
     return user
 
 
+def get_redirect_client_url(request):
+    bearer_token, date_expired = request.user.create_bearer_token(request, age=3600*36*5)
+    data = {
+        'type': 'auth',
+        'bearer_token': bearer_token,
+        'date_expired': date_expired.timestamp()
+    }
+    buf = base64.b64encode(json.dumps(data).encode()).decode()
+    redirect_url = 'jms://{}'.format(buf)
+    message_data = {
+        'title': _('Auth success'),
+        'message': _("Redirecting to JumpServer Client"),
+        'redirect_url': redirect_url,
+        'interval': 1,
+        'has_cancel': False,
+    }
+    url = FlashMessageUtil.gen_message_url(message_data)
+    return url
+
+
 def redirect_user_first_login_or_index(request, redirect_field_name):
     url = request.POST.get(redirect_field_name)
     if not url:
         url = request.GET.get(redirect_field_name)
+
+    if url == 'client':
+        url = get_redirect_client_url(request)
+
     url = safe_next_url(url, request=request)
     # 防止 next 地址为 None
     if not url or url.lower() in ['none']: