From 9b00375d93bb5f028018e75239453c2c53da74d0 Mon Sep 17 00:00:00 2001
From: xgdyp <zxyasrf@163.com>
Date: Fri, 9 Jul 2021 22:37:20 +0800
Subject: [PATCH] restrict file type

---
 be/app/api/v1/files.py   |  5 +++--
 be/app/api/v1/project.py | 12 +++++++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/be/app/api/v1/files.py b/be/app/api/v1/files.py
index af9940c..9231704 100644
--- a/be/app/api/v1/files.py
+++ b/be/app/api/v1/files.py
@@ -66,7 +66,8 @@ def query_file():
 def get_json():
     ret_info = ReturnInfo()
     try:
-        project_name = request.args.get("projectName").strip()
+        
+        project_name = request.args.get("projectName")
         download_json = []
         anno_data = read_json_file(PROJECT_PATH.format(project_name)+'/anno.json')
         for item in anno_data:
@@ -100,7 +101,7 @@ def get_json():
     #使用response可以将result.json再删除掉
     response = make_response(send_from_directory(directory='', path=DOWNLOAD_FILE_LOCATION.format(project_name),
                                                  as_attachment=True))
-    response.headers["Content-disposition"] = 'attachment; filename={}_result.json'.format(project_name)
+    response.headers["Content-disposition"] = 'attachment; filename=result.json'
     # print(PROJECT_PATH.format(project_name)+'/result.json')
     # os.remove(PROJECT_PATH.format(project_name)+'/result.json')
 
diff --git a/be/app/api/v1/project.py b/be/app/api/v1/project.py
index 071468c..8108f6c 100644
--- a/be/app/api/v1/project.py
+++ b/be/app/api/v1/project.py
@@ -34,15 +34,17 @@ def get_zipped_data():
 
         # Put files in the folder directly under the project directory
         for item in os.listdir(target_path):
+
             # If item is a folder,copy files in item to project folder,then remove this folder
             folder_path = target_path + '/' + item
             if os.path.isdir(folder_path):
                 for file in os.listdir(folder_path):
-                    shutil.copy(folder_path + '/' + file, folder_path + '/../')
-                    # Here  to deal with chinese encode in module zipfile and rarfile
-                    if file_type == 'zip':
-                        os.rename(folder_path + '/../' + file,
-                                  folder_path + '/../' + file.encode('cp437').decode('GBK'))
+                    if file.split('.')[-1] == 'txt':
+                        shutil.copy(folder_path + '/' + file, folder_path + '/../')
+                        # Here  to deal with chinese encode in module zipfile and rarfile
+                        if file_type == 'zip':
+                            os.rename(folder_path + '/../' + file,
+                                      folder_path + '/../' + file.encode('cp437').decode('GBK'))
                 shutil.rmtree(folder_path)
     except Exception as e:
         print(e)