feat: 录像增加 cast 格式 (#7259)

* feat: 录像增加 cast 格式

* perf: 优化一下写法

* perf: 修改一点点

Co-authored-by: ibuler <ibuler@qq.com>

apps/terminal/api/session.py

@@ -28,7 +28,6 @@ from ..hands import SystemUser
 from ..models import Session
 from .. import serializers
 
-
 __all__ = [
     'SessionViewSet', 'SessionReplayViewSet', 'SessionJoinValidateAPI'
 ]
@@ -41,7 +40,7 @@ class SessionViewSet(OrgBulkModelViewSet):
         'default': serializers.SessionSerializer,
         'display': serializers.SessionDisplaySerializer,
     }
-    permission_classes = (IsOrgAdminOrAppUser, )
+    permission_classes = (IsOrgAdminOrAppUser,)
     search_fields = [
         "user", "asset", "system_user", "remote_addr", "protocol", "is_finished", 'login_from',
     ]
@@ -71,7 +70,8 @@ class SessionViewSet(OrgBulkModelViewSet):
         os.chdir(current_dir)
         return file
 
-    @action(methods=[GET], detail=True, renderer_classes=(PassthroughRenderer,), url_path='replay/download', url_name='replay-download')
+    @action(methods=[GET], detail=True, renderer_classes=(PassthroughRenderer,), url_path='replay/download',
+            url_name='replay-download')
     def download(self, request, *args, **kwargs):
         session = self.get_object()
         local_path, url = utils.get_session_replay_url(session)
@@ -102,7 +102,7 @@ class SessionViewSet(OrgBulkModelViewSet):
 
     def get_permissions(self):
         if self.request.method.lower() in ['get', 'options']:
-            self.permission_classes = (IsOrgAdminOrAppUser | IsOrgAuditor, )
+            self.permission_classes = (IsOrgAdminOrAppUser | IsOrgAuditor,)
         return super().get_permissions()
 
 
@@ -119,7 +119,9 @@ class SessionReplayViewSet(AsyncApiMixin, viewsets.ViewSet):
 
         if serializer.is_valid():
             file = serializer.validated_data['file']
-            name, err = session.save_replay_to_storage(file)
+            # 兼容旧版本 API 未指定 version 为 2 的情况
+            version = serializer.validated_data.get('version', 2)
+            name, err = session.save_replay_to_storage_with_version(file, version)
             if not name:
                 msg = "Failed save replay `{}`: {}".format(session_id, err)
                 logger.error(msg)
@@ -137,6 +139,8 @@ class SessionReplayViewSet(AsyncApiMixin, viewsets.ViewSet):
         if session.protocol in ('rdp', 'vnc'):
             # 需要考虑录像播放和离线播放器的约定，暂时不处理
             tp = 'guacamole'
+        if url.endswith('.cast.gz'):
+            tp = 'asciicast'
 
         download_url = reverse('api-terminal:session-replay-download', kwargs={'pk': session.id})
         data = {
@@ -168,7 +172,7 @@ class SessionReplayViewSet(AsyncApiMixin, viewsets.ViewSet):
 
 
 class SessionJoinValidateAPI(views.APIView):
-    permission_classes = (IsAppUser, )
+    permission_classes = (IsAppUser,)
     serializer_class = serializers.SessionJoinValidateSerializer
 
     def post(self, request, *args, **kwargs):

apps/terminal/models/session.py

@@ -56,26 +56,65 @@ class Session(OrgModelMixin):
     upload_to = 'replay'
     ACTIVE_CACHE_KEY_PREFIX = 'SESSION_ACTIVE_{}'
     _DATE_START_FIRST_HAS_REPLAY_RDP_SESSION = None
+    SUFFIX_MAP = {1: '.gz', 2: '.replay.gz', 3: '.cast.gz'}
+    DEFAULT_SUFFIXES = ['.replay.gz', '.cast.gz', '.gz']
 
-    def get_rel_replay_path(self, version=2):
+    # Todo: 将来干掉 local_path, 使用 default storage 实现
+    def get_all_possible_local_path(self):
         """
-        获取session日志的文件路径
-        :param version: 原来后缀是 .gz，为了统一新版本改为 .replay.gz
+        获取所有可能的本地存储录像文件路径
+        :return:
+        """
+        return [self.get_local_storage_path_by_suffix(suffix)
+                for suffix in self.SUFFIX_MAP.values()]
+
+    def get_all_possible_relative_path(self):
+        """
+        获取所有可能的外部存储录像文件路径
+        :return:
+        """
+        return [self.get_relative_path_by_suffix(suffix)
+                for suffix in self.SUFFIX_MAP.values()]
+
+    def get_local_storage_path_by_suffix(self, suffix='.cast.gz'):
+        """
+        local_path: replay/2021-12-08/session_id.cast.gz
+        通过后缀名获取本地存储的录像文件路径
+        :param suffix: .cast.gz | '.replay.gz' | '.gz'
+        :return:
+        """
+        rel_path = self.get_relative_path_by_suffix(suffix)
+        if suffix == '.gz':
+            # 兼容 v1 的版本
+            return rel_path
+        return os.path.join(self.upload_to, rel_path)
+
+    def get_relative_path_by_suffix(self, suffix='.cast.gz'):
+        """
+        relative_path: 2021-12-08/session_id.cast.gz
+        通过后缀名获取外部存储录像文件路径
+        :param suffix: .cast.gz | '.replay.gz' | '.gz'
         :return:
         """
-        suffix = '.replay.gz'
-        if version == 1:
-            suffix = '.gz'
         date = self.date_start.strftime('%Y-%m-%d')
         return os.path.join(date, str(self.id) + suffix)
 
-    def get_local_path(self, version=2):
-        rel_path = self.get_rel_replay_path(version=version)
-        if version == 2:
-            local_path = os.path.join(self.upload_to, rel_path)
-        else:
-            local_path = rel_path
-        return local_path
+    def get_local_path_by_relative_path(self, rel_path):
+        """
+        2021-12-08/session_id.cast.gz
+        :param rel_path:
+        :return: replay/2021-12-08/session_id.cast.gz
+        """
+        return '{}/{}'.format(self.upload_to, rel_path)
+
+    def get_relative_path_by_local_path(self, local_path):
+        return local_path.replace('{}/'.format(self.upload_to), '')
+
+    def find_ok_relative_path_in_storage(self, storage):
+        session_paths = self.get_all_possible_relative_path()
+        for rel_path in session_paths:
+            if storage.exists(rel_path):
+                return rel_path
 
     @property
     def asset_obj(self):
@@ -133,8 +172,9 @@ class Session(OrgModelMixin):
         else:
             return True
 
-    def save_replay_to_storage(self, f):
-        local_path = self.get_local_path()
+    def save_replay_to_storage_with_version(self, f, version=2):
+        suffix = self.SUFFIX_MAP.get(version, '.cast.gz')
+        local_path = self.get_local_storage_path_by_suffix(suffix)
         try:
             name = default_storage.save(local_path, f)
         except OSError as e:
@@ -148,7 +188,7 @@ class Session(OrgModelMixin):
     @classmethod
     def set_sessions_active(cls, session_ids):
         data = {cls.ACTIVE_CACHE_KEY_PREFIX.format(i): i for i in session_ids}
-        cache.set_many(data, timeout=5*60)
+        cache.set_many(data, timeout=5 * 60)
 
     @classmethod
     def get_active_sessions(cls):
@@ -195,7 +235,7 @@ class Session(OrgModelMixin):
             for user, asset, system_user in ziped:
                 ip = random_ip()
                 date_start = random_datetime(month_ago, now)
-                date_end = random_datetime(date_start, date_start+timezone.timedelta(hours=2))
+                date_end = random_datetime(date_start, date_start + timezone.timedelta(hours=2))
                 data = dict(
                     user=str(user), user_id=user.id,
                     asset=str(asset), asset_id=asset.id,

apps/terminal/serializers/session.py

@@ -50,6 +50,7 @@ class SessionDisplaySerializer(SessionSerializer):
 
 class ReplaySerializer(serializers.Serializer):
     file = serializers.FileField(allow_empty_file=True)
+    version = serializers.IntegerField(write_only=True, required=False, min_value=2, max_value=3)
 
 
 class SessionJoinValidateSerializer(serializers.Serializer):

apps/terminal/tasks.py

@@ -85,7 +85,7 @@ def upload_session_replay_to_external_storage(session_id):
         logger.error(f'Session replay not found, may be upload error: {local_path}')
         return
     abs_path = default_storage.path(local_path)
-    remote_path = session.get_rel_replay_path()
+    remote_path = session.get_relative_path_by_local_path(abs_path)
     ok, err = server_replay_storage.upload(abs_path, remote_path)
     if not ok:
         logger.error(f'Session replay upload to external error: {err}')

apps/terminal/utils.py

@@ -1,30 +1,28 @@
 # -*- coding: utf-8 -*-
 #
 import os
-from itertools import groupby
+from itertools import groupby, chain
 
 from django.conf import settings
 from django.core.files.storage import default_storage
-from django.utils.translation import ugettext as _
 
 import jms_storage
 
-from common.tasks import send_mail_async
-from common.utils import get_logger, reverse
+from common.utils import get_logger
 from . import const
-from .models import ReplayStorage, Session, Command
+from .models import ReplayStorage
 
 logger = get_logger(__name__)
 
 
 def find_session_replay_local(session):
-    # 新版本和老版本的文件后缀不同
-    session_path = session.get_rel_replay_path()  # 存在外部存储上的路径
-    local_path = session.get_local_path()
-    local_path_v1 = session.get_local_path(version=1)
+    # 存在外部存储上，所有可能的路径名
+    session_paths = session.get_all_possible_relative_path()
 
-    # 去default storage中查找
-    for _local_path in (local_path, local_path_v1, session_path):
+    # 存在本地存储上，所有可能的路径名
+    local_paths = session.get_all_possible_local_path()
+
+    for _local_path in chain(session_paths, local_paths):
         if default_storage.exists(_local_path):
             url = default_storage.url(_local_path)
             return _local_path, url
@@ -32,8 +30,6 @@ def find_session_replay_local(session):
 
 
 def download_session_replay(session):
-    session_path = session.get_rel_replay_path()  # 存在外部存储上的路径
-    local_path = session.get_local_path()
     replay_storages = ReplayStorage.objects.all()
     configs = {
         storage.name: storage.config
@@ -45,13 +41,23 @@ def download_session_replay(session):
     if not configs:
         msg = "Not found replay file, and not remote storage set"
         return None, msg
+    storage = jms_storage.get_multi_object_storage(configs)
+
+    # 获取外部存储路径名
+    session_path = session.find_ok_relative_path_in_storage(storage)
+    if not session_path:
+        msg = "Not found session replay file"
+        return None, msg
+
+    # 通过外部存储路径名后缀，构造真实的本地存储路径
+    local_path = session.get_local_path_by_relative_path(session_path)
 
     # 保存到storage的路径
     target_path = os.path.join(default_storage.base_location, local_path)
     target_dir = os.path.dirname(target_path)
     if not os.path.isdir(target_dir):
         os.makedirs(target_dir, exist_ok=True)
-    storage = jms_storage.get_multi_object_storage(configs)
+
     ok, err = storage.download(session_path, target_path)
     if not ok:
         msg = "Failed download replay file: {}".format(err)