From b92ccdd05eb9e7e753c1548c8dfbb88d7eef4f0a Mon Sep 17 00:00:00 2001
From: liuzheng <liuzheng712@gmail.com>
Date: Mon, 12 Mar 2018 10:30:40 +0800
Subject: [PATCH] fix: update the jms-storage==0.0.11

---
 apps/terminal/api.py          | 44 +++++++++++------------------------
 requirements/requirements.txt |  1 +
 2 files changed, 14 insertions(+), 31 deletions(-)

diff --git a/apps/terminal/api.py b/apps/terminal/api.py
index 370835ca2..1f113ca85 100644
--- a/apps/terminal/api.py
+++ b/apps/terminal/api.py
@@ -4,7 +4,6 @@ from collections import OrderedDict
 import logging
 import os
 import uuid
-import boto3  # AWS S3 sdk
 
 from django.core.cache import cache
 from django.shortcuts import get_object_or_404, redirect
@@ -13,6 +12,8 @@ from django.core.files.storage import default_storage
 from django.http import HttpResponseNotFound
 from django.conf import settings
 
+import jms_storage
+
 from rest_framework import viewsets, serializers
 from rest_framework.views import APIView, Response
 from rest_framework.permissions import AllowAny
@@ -282,36 +283,17 @@ class SessionReplayViewSet(viewsets.ViewSet):
             url = default_storage.url(path)
             return redirect(url)
         else:
-            config = settings.TERMINAL_REPLAY_STORAGE.items()
-            if config:
-                for name, value in config:
-                    if value.get("TYPE", '') == "s3":
-                        client, bucket = self.s3Client(value)
-                        try:
-                            date = self.session.date_start.strftime('%Y-%m-%d')
-
-                            client.head_object(Bucket=bucket,
-                                               Key=os.path.join(date, str(self.session.id) + '.replay.gz'))
-                            client.download_file(bucket, os.path.join(date, str(self.session.id) + '.replay.gz'),
-                                                 default_storage.base_location + '/' + path)
-                            return redirect(default_storage.url(path))
-                        except:
-                            pass
-            return HttpResponseNotFound()
-
-    def s3Client(self, config):
-        bucket = config.get("BUCKET", "jumpserver")
-        REGION = config.get("REGION", None)
-        ACCESS_KEY = config.get("ACCESS_KEY", None)
-        SECRET_KEY = config.get("SECRET_KEY", None)
-        if ACCESS_KEY and REGION and SECRET_KEY:
-            s3 = boto3.client('s3',
-                              region_name=REGION,
-                              aws_access_key_id=ACCESS_KEY,
-                              aws_secret_access_key=SECRET_KEY)
-        else:
-            s3 = boto3.client('s3')
-        return s3, bucket
+            configs = settings.TERMINAL_REPLAY_STORAGE.items()
+            if configs:
+                for name, config in configs:
+                    client = jms_storage.init(config)
+                    date = self.session.date_start.strftime('%Y-%m-%d')
+                    if client.has_file(os.path.join(date, str(self.session.id) + '.replay.gz')) \
+                            and \
+                            client.download_file(os.path.join(date, str(self.session.id) + '.replay.gz'),
+                                                 default_storage.base_location + '/' + path):
+                        return redirect(default_storage.url(path))
+        return HttpResponseNotFound()
 
 
 class TerminalConfig(APIView):
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 200d23480..536d00bc6 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -61,6 +61,7 @@ pytz==2017.3
 PyYAML==3.12
 redis==2.10.6
 requests==2.18.4
+jms-storage==0.0.11
 s3transfer==0.1.13
 simplejson==3.13.2
 six==1.11.0