1
#!/usr/bin/env python3
2

3
# Contest Management System - http://cms-dev.github.io/
4
# Copyright © 2014-2015 William Di Luigi <williamdiluigi@gmail.com>
5
# Copyright © 2014 Stefano Maggiolo <s.maggiolo@gmail.com>
6
#
7
# This program is free software: you can redistribute it and/or modify
8
# it under the terms of the GNU Affero General Public License as
9
# published by the Free Software Foundation, either version 3 of the
10
# License, or (at your option) any later version.
11
#
12
# This program is distributed in the hope that it will be useful,
13
# but WITHOUT ANY WARRANTY; without even the implied warranty of
14
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15
# GNU Affero General Public License for more details.
16
#
17
# You should have received a copy of the GNU Affero General Public License
18
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
19

20 2
"""Abstraction layer for reading from and writing to archives.
21

22
"""
23

24 2
import os
25 2
import shutil
26 2
import tempfile
27

28 2
import patoolib
29 2
from patoolib.util import PatoolError
30

31 2
from cms import config
32

33

34 2
class ArchiveException(Exception):
35
    """Exception for when the interaction with the Archive class is
36
    incorrect.
37

38
    """
39 2
    pass
40

41

42 2
class Archive:
43
    """Class to manage archives.
44

45
    This class has static methods to test, extract, and create
46
    archives. Moreover, an instance of this class can be create to
47
    manage an existing archive. At the moment, all operations depend
48
    on calling first the unpack method, that extract the archive in a
49
    temporary directory.
50

51
    """
52

53 2
    @staticmethod
54
    def is_supported(path):
55
        """Return whether the file at path is supported by patoolib.
56

57
        path (string): the path to test.
58

59
        return (bool): whether path is supported.
60

61
        """
62 1
        try:
63 1
            patoolib.test_archive(path, interactive=False)
64 1
            return True
65 1
        except PatoolError:
66 1
            return False
67

68 2
    @staticmethod
69
    def create_from_dir(from_dir, archive_path):
70
        """Create a new archive containing all files in from_dir.
71

72
        from_dir (string): directory with the files to archive.
73
        archive_path (string): the new archive's path.
74

75
        """
76 0
        files = tuple(os.listdir(from_dir))
77 0
        cwd = os.getcwd()
78 0
        os.chdir(from_dir)
79 0
        patoolib.create_archive(archive_path, files, interactive=False)
80 0
        os.chdir(cwd)
81

82 2
    @staticmethod
83
    def extract_to_dir(archive_path, to_dir):
84
        """Extract the content of an archive in to_dir.
85

86
        archive_path (string): path of the archive to extract.
87
        to_dir (string): destination directory.
88

89
        """
90 0
        patoolib.extract_archive(archive_path, outdir=to_dir, interactive=False)
91

92 2
    @staticmethod
93
    def from_raw_data(raw_data):
94
        """Create an Archive object out of raw archive data.
95

96
        This method treats the given string as archive data: it dumps it
97
        into a temporary file, then creates an Archive object. Since the
98
        user did not provide a path, we assume that when cleanup() is
99
        called the temporary file should be deleted as well as unpacked
100
        data.
101

102
        raw_data (bytes): the actual bytes that form the archive.
103

104
        return (Archive|None): an object that represents the new
105
            archive or None, if raw_data doesn't represent an archive.
106

107
        """
108 1
        temp_file, temp_filename = tempfile.mkstemp(dir=config.temp_dir)
109 1
        with open(temp_file, "wb") as temp_file:
110 1
            temp_file.write(raw_data)
111

112 1
        try:
113 1
            return Archive(temp_filename, delete_source=True)
114 1
        except ArchiveException:
115 1
            os.remove(temp_filename)
116 1
            return None
117

118 2
    def __init__(self, path, delete_source=False):
119
        """Init.
120

121
        path (string): the path of the archive.
122
        delete_source (bool): whether the source archive should be
123
            deleted at cleanup or not.
124

125
        """
126 1
        if not Archive.is_supported(path):
127 1
            raise ArchiveException("This type of archive is not supported.")
128 1
        self.delete_source = delete_source
129 1
        self.path = path
130 1
        self.temp_dir = None
131

132 2
    def unpack(self):
133
        """Extract archive's content to a temporary directory.
134

135
        return (string): the path of the temporary directory.
136

137
        """
138 1
        self.temp_dir = tempfile.mkdtemp(dir=config.temp_dir)
139 1
        patoolib.extract_archive(self.path, outdir=self.temp_dir,
140
                                 interactive=False)
141 1
        return self.temp_dir
142

143 2
    def repack(self, target):
144
        """Repack to a new archive all the files which were unpacked in
145
        self.temp_dir.
146

147
        target (string): the new archive path.
148

149
        """
150 0
        if self.temp_dir is None:
151 0
            raise ArchiveException("The unpack() method must be called first.")
152 0
        Archive.create_from_dir(self.temp_dir, target)
153

154 2
    def cleanup(self):
155
        """Remove temporary directory, if needed.
156

157
        """
158 1
        if self.temp_dir is not None and os.path.exists(self.temp_dir):
159 1
            shutil.rmtree(self.temp_dir)
160 1
            self.temp_dir = None
161 1
        if self.delete_source:
162 1
            try:
163 1
                os.remove(self.path)
164 0
            except OSError:
165
                # Cannot delete source, it is not a big problem.
166 0
                pass
167

168 2
    def namelist(self):
169
        """Returns all pathnames for this archive.
170

171
        return ([string]): list of files in the archive.
172

173
        raise (NotImplementedError): when the archive was unpacked
174
            first.
175

176
        """
177 1
        if self.temp_dir is None:
178
            # Unfortunately, this "prints" names to the screen, so it's
179
            # not very handy.
180
            # patoolib.list_archive(self.path)
181 0
            raise NotImplementedError("Cannot list before unpacking.")
182
        else:
183 1
            names = []
184 1
            for path, _, filenames in os.walk(self.temp_dir):
185 1
                for filename in filenames:
186 1
                    names.append(os.path.relpath(os.path.join(path, filename),
187
                                                 self.temp_dir))
188 1
            return names
189

190 2
    def read(self, file_path):
191
        """Read a single file and return its file object.
192

193
        file_path (string): path of the file in the archive.
194

195
        return (file): handler for the file.
196

197
        raise (NotImplementedError): when the archive was unpacked
198
            first.
199

200
        """
201 1
        if self.temp_dir is None:
202
            # Unfortunately, patoolib does not expose an API to do this.
203 0
            raise NotImplementedError("Cannot read before unpacking.")
204
        else:
205 1
            return open(os.path.join(self.temp_dir, file_path), "rb")
206

207 2
    def write(self, file_path, file_object):
208
        """Writes a file in the archive in place.
209

210
        file_path (string): new path in the archive.
211
        file_object (object): file-like object.
212

213
        raise (NotImplementedError): always; this method is not yet
214
            implemented.
215

216
        """
217 0
        if self.temp_dir is None:
218
            # Unfortunately, patoolib does not expose an API to do this.
219 0
            raise NotImplementedError("Cannot write before unpacking.")
220
        else:
221 0
            raise NotImplementedError(
222
                "You should write the file directly, in the "
223
                "folder returned by unpack(), and then "
224
                "call the repack() method.")

Read our documentation on viewing source code .

Loading