From 6b6900cfd4c16d6095db2c744e7b100cbc5a415b Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Mon, 8 May 2023 17:02:55 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E5=AF=BC=E5=87=BAaccount=20=E5=88=97?=
 =?UTF-8?q?=E8=A1=A8=E9=9C=80=E5=AF=B9=E6=96=87=E4=BB=B6=E8=BF=9B=E8=A1=8C?=
 =?UTF-8?q?=E5=8A=A0=E5=AF=86=20(#10366)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 apps/common/drf/renders/base.py      | 36 +++++++++++++++++++++++++++-
 apps/locale/ja/LC_MESSAGES/django.po |  6 +++++
 apps/locale/zh/LC_MESSAGES/django.po |  7 ++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/apps/common/drf/renders/base.py b/apps/common/drf/renders/base.py
index 5b2aa5612..cc92eb8c2 100644
--- a/apps/common/drf/renders/base.py
+++ b/apps/common/drf/renders/base.py
@@ -1,6 +1,10 @@
 import abc
+import io
+import re
 from datetime import datetime
 
+import pyzipper
+from django.utils.translation import ugettext_lazy as _
 from rest_framework import serializers
 from rest_framework.renderers import BaseRenderer
 from rest_framework.utils import encoders, json
@@ -180,9 +184,39 @@ class BaseFileRenderer(BaseRenderer):
             self.write_column_titles(column_titles)
             self.write_rows(rows)
             self.after_render()
-            value = self.get_rendered_value()
+            value = self.compress_into_zip_file(view, request, response)
         except Exception as e:
             logger.debug(e, exc_info=True)
             value = 'Render error! ({})'.format(self.media_type).encode('utf-8')
             return value
         return value
+
+    def compress_into_zip_file(self, view, request, response):
+        value = self.get_rendered_value()
+        from accounts.models import Account
+        if str(view.model) not in (str(Account),) or self.template != 'export':
+            return value
+
+        filename_pattern = re.compile(r'filename="([^"]+)"')
+        content_disposition = response['Content-Disposition']
+        match = filename_pattern.search(content_disposition)
+        filename = match.group(1)
+        response['Content-Disposition'] = content_disposition.replace(self.format, 'zip')
+
+        contents_io = io.BytesIO()
+        secret_key = request.user.secret_key
+        if not secret_key:
+            content = _("{} - The encryption password has not been set - "
+                        "please go to personal information -> file encryption password "
+                        "to set the encryption password").format(request.user.name)
+
+            response['Content-Disposition'] = content_disposition.replace(self.format, 'txt')
+            contents_io.write(content.encode('utf-8'))
+            return contents_io.getvalue()
+
+        with pyzipper.AESZipFile(
+                contents_io, 'w', compression=pyzipper.ZIP_LZMA, encryption=pyzipper.WZ_AES
+        ) as zf:
+            zf.setpassword(secret_key.encode('utf8'))
+            zf.writestr(filename, value)
+        return contents_io.getvalue()
diff --git a/apps/locale/ja/LC_MESSAGES/django.po b/apps/locale/ja/LC_MESSAGES/django.po
index b751ef764..759def5fd 100644
--- a/apps/locale/ja/LC_MESSAGES/django.po
+++ b/apps/locale/ja/LC_MESSAGES/django.po
@@ -547,6 +547,12 @@ msgstr ""
 "{} -暗号化変更タスクが完了しました: 暗号化パスワードが設定されていません-個人"
 "情報にアクセスしてください-> ファイル暗号化パスワードを設定してください"
 
+msgid ""
+"{} - The encryption password has not been set - please go to personal information -> file encryption "
+"password to set the encryption password"
+msgstr ""
+"{} - 暗号化パスワードが設定されていません-個人情報->ファイル暗号化パスワードに暗号化パスワードを設定してください"
+
 #: accounts/serializers/account/account.py:28
 msgid "Push now"
 msgstr "今すぐプッシュ"
diff --git a/apps/locale/zh/LC_MESSAGES/django.po b/apps/locale/zh/LC_MESSAGES/django.po
index 40199acec..ff2b9c6a2 100644
--- a/apps/locale/zh/LC_MESSAGES/django.po
+++ b/apps/locale/zh/LC_MESSAGES/django.po
@@ -524,6 +524,13 @@ msgstr ""
 "{} - 账号备份任务已完成: 未设置加密密码 - 请前往个人信息 -> 文件加密密码中设"
 "置加密密码"
 
+msgid ""
+"{} - The encryption password has not been set - please go to personal information -> file "
+"encryption password to set the encryption password"
+msgstr ""
+"{} - 未设置加密密码 - 请前往个人信息 -> 文件加密密码中设"
+"置加密密码"
+
 #: accounts/notifications.py:33
 msgid "Notification of implementation result of encryption change plan"
 msgstr "改密计划任务结果通知"