From ff00d880a0718ae10f06243dc0b341b7fa3a6743 Mon Sep 17 00:00:00 2001
From: spxcds <jvbei@live.com>
Date: Wed, 6 Apr 2016 16:18:31 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=B5=8B=E8=AF=95=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E4=B8=8B=E8=BD=BD=E7=9A=84=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore       |  3 ++-
 oj/urls.py       |  3 ++-
 problem/views.py | 59 +++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1be93cc8..9cb7a002 100644
--- a/.gitignore
+++ b/.gitignore
@@ -65,4 +65,5 @@ test_case/
 release/
 upload/
 custom_settings.py
-docker-compose.yml
\ No newline at end of file
+docker-compose.yml
+*.zip
diff --git a/oj/urls.py b/oj/urls.py
index f8625825..0cd8b225 100644
--- a/oj/urls.py
+++ b/oj/urls.py
@@ -20,7 +20,7 @@ from group.views import (GroupAdminAPIView, GroupMemberAdminAPIView,
 
 from admin.views import AdminTemplateView
 
-from problem.views import TestCaseUploadAPIView, ProblemTagAdminAPIView, ProblemAdminAPIView, OpenAPIProblemAPI
+from problem.views import TestCaseUploadAPIView, TestCaseDownloadAPIView, ProblemTagAdminAPIView, ProblemAdminAPIView, OpenAPIProblemAPI
 from submission.views import (SubmissionAPIView, SubmissionAdminAPIView, ContestSubmissionAPIView,
                               SubmissionShareAPIView, SubmissionRejudgeAdminAPIView, OpenAPISubmitCodeAPI)
 from judge_dispatcher.views import AdminJudgeServerAPIView
@@ -68,6 +68,7 @@ urlpatterns = [
     url(r'^api/admin/contest_problem/public/', MakeContestProblemPublicAPIView.as_view(),
         name="make_contest_problem_public"),
     url(r'^api/admin/test_case_upload/$', TestCaseUploadAPIView.as_view(), name="test_case_upload_api"),
+    url(r'^api/admin/test_case_download/$', TestCaseDownloadAPIView.as_view(), name="test_case_download_api"),
     url(r'^api/admin/tag/$', ProblemTagAdminAPIView.as_view(), name="problem_tag_admin_api"),
     url(r'^api/admin/join_group_request/$', JoinGroupRequestAdminAPIView.as_view(),
         name="join_group_request_admin_api"),
diff --git a/problem/views.py b/problem/views.py
index 27f8c705..80bb6dd6 100644
--- a/problem/views.py
+++ b/problem/views.py
@@ -6,6 +6,7 @@ import hashlib
 import json
 import logging
 
+from django.http import StreamingHttpResponse
 from django.shortcuts import render
 from django.db.models import Q, Count
 from django.core.paginator import Paginator
@@ -14,7 +15,8 @@ from django.conf import settings
 from rest_framework.views import APIView
 
 from account.models import SUPER_ADMIN, User
-from account.decorators import super_admin_required
+from account.decorators import super_admin_required, admin_required
+from contest.models import ContestProblem
 from utils.shortcuts import (serializer_invalid_response, error_response,
                              success_response, paginate, rand_str, error_page)
 from .serizalizers import (CreateProblemSerializer, EditProblemSerializer, ProblemSerializer,
@@ -282,6 +284,61 @@ class TestCaseUploadAPIView(APIView):
         return success_response({"file_list": config["test_cases"]})
 
 
+class TestCaseDownloadAPIView(APIView):
+    """
+    下载题目的测试数据
+    """
+
+    def _is_legal_test_case_file_name(self, file_name):
+        regex = r"^[1-9]\d*\.(in|out)$"
+        return re.compile(regex).match(file_name) is not None
+
+    def file_iterator(self, big_file, chunk_size=512):
+        with open(big_file) as f:
+            while True:
+                c = f.read(chunk_size)
+                if c:
+                    yield c
+                else:
+                    break
+
+    @admin_required
+    def get(self, request):
+        test_case_id = request.GET.get("test_case_id", None)
+        if not test_case_id:
+            return error_response(u"参数错误")
+        # 防止URL./../../.上层目录遍历
+        if not re.compile(r"^[1-9a-zA-Z]+$").match(test_case_id):
+            return error_response(u"参数错误")
+
+        try:
+            # 超级管理员可以下载全部的题目的测试数据
+            # 普通管理员只能下载自己创建的题目的测试数据
+            if request.user.admin_type != SUPER_ADMIN:
+                ContestProblem.objects.get(test_case_id=test_case_id, created_by=request.user)
+
+            test_case_dir = os.path.join(settings.TEST_CASE_DIR, test_case_id)
+            if not os.path.exists(test_case_dir):
+                return error_response(u"测试用例不存在")
+
+            # 压缩测试用例,命名规则为 "test_case" + test_case_id + ".zip"
+            test_case_zip = os.path.join("/tmp", "test_case-" + test_case_id + ".zip")
+
+            zf = zipfile.ZipFile(test_case_zip, "w", zipfile.ZIP_DEFLATED)
+            for filename in os.listdir(test_case_dir):
+                if self._is_legal_test_case_file_name(filename):
+                    zf.write(os.path.join(test_case_dir, filename), filename)
+            zf.close()
+
+            # 大文件传输
+            response = StreamingHttpResponse(self.file_iterator(test_case_zip))
+            response['Content-Type'] = 'application/octet-stream'
+            response['Content-Disposition'] = 'attachment;filename=test_case-%s.zip' % test_case_id
+            return response
+        except ContestProblem.DoesNotExist:
+            return error_response(u"题目不存在")
+
+
 def problem_list_page(request, page=1):
     """
     前台的问题列表