CWS: use cmscommon.archive when suitable

Author: wil93

cms/server/ContestWebServer.py

@@ -9,6 +9,7 @@
 # Copyright © 2013 Bernard Blackham <bernard@largestprime.net>
 # Copyright © 2014 Artem Iglikov <artem.iglikov@gmail.com>
 # Copyright © 2014 Fabian Gundlach <320pointsguy@gmail.com>
+# Copyright © 2015 William Di Luigi <williamdiluigi@gmail.com>
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -53,7 +54,6 @@
 import re
 import socket
 import struct
-import tempfile
 import traceback
 from datetime import timedelta
 from urllib import quote
@@ -72,15 +72,16 @@
 from cms.db.filecacher import FileCacher
 from cms.grading.tasktypes import get_task_type
 from cms.grading.scoretypes import get_score_type
-from cms.server import file_handler_gen, extract_archive, \
-    actual_phase_required, get_url_root, filter_ascii, \
-    CommonRequestHandler, format_size, compute_actual_phase
+from cms.server import file_handler_gen, actual_phase_required, \
+    get_url_root, filter_ascii, CommonRequestHandler, format_size, \
+    compute_actual_phase
 from cmscommon.isocodes import is_language_code, translate_language_code, \
     is_country_code, translate_country_code, \
     is_language_country_code, translate_language_country_code
 from cmscommon.crypto import encrypt_number
 from cmscommon.datetime import make_datetime, make_timestamp, get_timezone
 from cmscommon.mimetypes import get_type_for_file_name
+from cmscommon.archive import Archive
 
 
 logger = logging.getLogger(__name__)
@@ -979,16 +980,10 @@ def post(self, task_name):
             archive_data = self.request.files["submission"][0]
             del self.request.files["submission"]
 
-            # Extract the files from the archive.
-            temp_archive_file, temp_archive_filename = \
-                tempfile.mkstemp(dir=config.temp_dir)
-            with os.fdopen(temp_archive_file, "w") as temp_archive_file:
-                temp_archive_file.write(archive_data["body"])
+            # Create the archive.
+            archive = Archive.from_raw_data(archive_data["body"])
 
-            archive_contents = extract_archive(temp_archive_filename,
-                                               archive_data["filename"])
-
-            if archive_contents is None:
+            if archive is None:
                 self.application.service.add_notification(
                     self.current_user.username,
                     self.timestamp,
@@ -999,8 +994,17 @@ def post(self, task_name):
                                                               safe=''))
                 return
 
-            for item in archive_contents:
-                self.request.files[item["filename"]] = [item]
+            # Extract the archive.
+            unpacked_dir = archive.unpack()
+            for name in archive.namelist():
+                filename = os.path.basename(name)
+                body = open(os.path.join(unpacked_dir, filename), "r").read()
+                self.request.files[filename] = [{
+                    'filename': filename,
+                    'body': body
+                }]
+
+            archive.cleanup()
 
         # This ensure that the user sent one file for every name in
         # submission format and no more. Less is acceptable if task
@@ -1518,16 +1522,10 @@ def post(self, task_name):
             archive_data = self.request.files["submission"][0]
             del self.request.files["submission"]
 
-            # Extract the files from the archive.
-            temp_archive_file, temp_archive_filename = \
-                tempfile.mkstemp(dir=config.temp_dir)
-            with os.fdopen(temp_archive_file, "w") as temp_archive_file:
-                temp_archive_file.write(archive_data["body"])
-
-            archive_contents = extract_archive(temp_archive_filename,
-                                               archive_data["filename"])
+            # Create the archive.
+            archive = Archive.from_raw_data(archive_data["body"])
 
-            if archive_contents is None:
+            if archive is None:
                 self.application.service.add_notification(
                     self.current_user.username,
                     self.timestamp,
@@ -1537,8 +1535,17 @@ def post(self, task_name):
                 self.redirect("/testing?%s" % quote(task.name, safe=''))
                 return
 
-            for item in archive_contents:
-                self.request.files[item["filename"]] = [item]
+            # Extract the archive.
+            unpacked_dir = archive.unpack()
+            for name in archive.namelist():
+                filename = os.path.basename(name)
+                body = open(os.path.join(unpacked_dir, filename), "r").read()
+                self.request.files[filename] = [{
+                    'filename': filename,
+                    'body': body
+                }]
+
+            archive.cleanup()
 
         # This ensure that the user sent one file for every name in
         # submission format and no more. Less is acceptable if task

cms/server/__init__.py

@@ -30,8 +30,6 @@
 
 import time
 import logging
-import tarfile
-import zipfile
 from datetime import datetime, timedelta
 from urllib import quote
 
@@ -199,49 +197,6 @@ def wrapped(self, *args, **kwargs):
     return decorator
 
 
-def extract_archive(temp_name, original_filename):
-    """Obtain a list of files inside the specified archive.
-
-    Returns a list of the files inside the archive located in
-    temp_name, using original_filename to guess the type of the
-    archive.
-
-    """
-    file_list = []
-    if original_filename.endswith(".zip"):
-        try:
-            zip_object = zipfile.ZipFile(temp_name, "r")
-            for item in zip_object.infolist():
-                file_list.append({
-                    "filename": item.filename,
-                    "body": zip_object.read(item)})
-        except Exception as error:
-            logger.warning("Exception while extracting zip file `%s'. %r",
-                           original_filename, error)
-            return None
-    elif original_filename.endswith(".tar.gz") \
-            or original_filename.endswith(".tar.bz2") \
-            or original_filename.endswith(".tar"):
-        try:
-            tar_object = tarfile.open(name=temp_name)
-            for item in tar_object.getmembers():
-                if item.isfile():
-                    file_list.append({
-                        "filename": item.name,
-                        "body": tar_object.extractfile(item).read()})
-        except tarfile.TarError:
-            logger.warning("Exception while extracting tar file `%s'. %r",
-                           original_filename, error)
-            return None
-        except IOError:
-            return None
-    else:
-        logger.warning("Compressed file `%s' not recognized.",
-                       original_filename)
-        return None
-    return file_list
-
-
 UNITS = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']
 DIMS = list(1024 ** x for x in xrange(9))
 

cmscommon/archive.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 # Contest Management System - http://cms-dev.github.io/
-# Copyright © 2014 William Di Luigi <williamdiluigi@gmail.com>
+# Copyright © 2014-2015 William Di Luigi <williamdiluigi@gmail.com>
 # Copyright © 2014 Stefano Maggiolo <s.maggiolo@gmail.com>
 #
 # This program is free software: you can redistribute it and/or modify
@@ -33,6 +33,8 @@
 
 from patoolib.util import PatoolError
 
+from cms import config
+
 
 class ArchiveException(Exception):
     """Exception for when the interaction with the Archive class is
@@ -92,14 +94,43 @@ def extract_to_dir(archive_path, to_dir):
         """
         patoolib.extract_archive(archive_path, outdir=to_dir)
 
-    def __init__(self, path):
+    @staticmethod
+    def from_raw_data(raw_data):
+        """Create an Archive object out of raw archive data.
+
+        This method treats the given string as archive data: it dumps it
+        into a temporary file, then creates an Archive object. Since the
+        user did not provide a path, we assume that when cleanup() is
+        called the temporary file should be deleted as well as unpacked
+        data.
+
+        raw_data (bytes): the actual bytes that form the archive.
+
+        return (Archive|None): an object that represents the new
+            archive or None, if raw_data doesn't represent an archive.
+
+        """
+        temp_file, temp_filename = tempfile.mkstemp(dir=config.temp_dir)
+        with os.fdopen(temp_file, "w") as temp_file:
+            temp_file.write(raw_data)
+
+        try:
+            return Archive(temp_filename, delete_source=True)
+        except ArchiveException:
+            os.remove(temp_filename)
+            return None
+
+    def __init__(self, path, delete_source=False):
         """Init.
 
         path (string): the path of the archive.
+        delete_source (bool): whether the source archive should be
+            deleted at cleanup or not.
 
         """
         if not Archive.is_supported(path):
             raise ArchiveException("This type of archive is not supported.")
+        self.delete_source = delete_source
         self.path = path
         self.temp_dir = None
 
@@ -109,7 +140,7 @@ def unpack(self):
         return (string): the path of the temporary directory.
 
         """
-        self.temp_dir = tempfile.mkdtemp()
+        self.temp_dir = tempfile.mkdtemp(dir=config.temp_dir)
         patoolib.extract_archive(self.path, outdir=self.temp_dir)
         return self.temp_dir
 
@@ -131,6 +162,12 @@ def cleanup(self):
         if os.path.exists(self.temp_dir):
             shutil.rmtree(self.temp_dir)
         self.temp_dir = None
+        if self.delete_source:
+            try:
+                os.remove(self.path)
+            except OSError:
+                # Cannot delete source, it is not a big problem.
+                pass
 
     def namelist(self):
         """Returns all pathnames for this archive.