1"""
2    salt.utils.gzip
3    ~~~~~~~~~~~~~~~
4    Helper module for handling gzip consistently between 2.7+ and 2.6-
5"""
6
7
8import gzip
9import io
10
11import salt.utils.files
12
13
14class GzipFile(gzip.GzipFile):
15    def __init__(self, filename=None, mode=None, compresslevel=9, fileobj=None):
16        gzip.GzipFile.__init__(self, filename, mode, compresslevel, fileobj)
17
18    ### Context manager (stolen from Python 2.7)###
19    def __enter__(self):
20        """Context management protocol.  Returns self."""
21        return self
22
23    def __exit__(self, *args):
24        """Context management protocol.  Calls close()"""
25        self.close()
26
27
28def open(filename, mode="rb", compresslevel=9):
29    if hasattr(gzip.GzipFile, "__enter__"):
30        return gzip.open(filename, mode, compresslevel)
31    else:
32        return GzipFile(filename, mode, compresslevel)
33
34
35def open_fileobj(fileobj, mode="rb", compresslevel=9):
36    if hasattr(gzip.GzipFile, "__enter__"):
37        return gzip.GzipFile(
38            filename="", mode=mode, fileobj=fileobj, compresslevel=compresslevel
39        )
40    return GzipFile(
41        filename="", mode=mode, fileobj=fileobj, compresslevel=compresslevel
42    )
43
44
45def compress(data, compresslevel=9):
46    """
47    Returns the data compressed at gzip level compression.
48    """
49    buf = io.BytesIO()
50    with open_fileobj(buf, "wb", compresslevel) as ogz:
51        if not isinstance(data, bytes):
52            data = data.encode(__salt_system_encoding__)
53        ogz.write(data)
54    compressed = buf.getvalue()
55    return compressed
56
57
58def uncompress(data):
59    buf = io.BytesIO(data)
60    with open_fileobj(buf, "rb") as igz:
61        unc = igz.read()
62        return unc
63
64
65def compress_file(fh_, compresslevel=9, chunk_size=1048576):
66    """
67    Generator that reads chunk_size bytes at a time from a file/filehandle and
68    yields the compressed result of each read.
69
70    .. note::
71        Each chunk is compressed separately. They cannot be stitched together
72        to form a compressed file. This function is designed to break up a file
73        into compressed chunks for transport and decompression/reassembly on a
74        remote host.
75    """
76    try:
77        bytes_read = int(chunk_size)
78        if bytes_read != chunk_size:
79            raise ValueError
80    except ValueError:
81        raise ValueError("chunk_size must be an integer")
82    try:
83        while bytes_read == chunk_size:
84            buf = io.BytesIO()
85            with open_fileobj(buf, "wb", compresslevel) as ogz:
86                try:
87                    bytes_read = ogz.write(fh_.read(chunk_size))
88                except AttributeError:
89                    # Open the file and re-attempt the read
90                    fh_ = salt.utils.files.fopen(fh_, "rb")
91                    bytes_read = ogz.write(fh_.read(chunk_size))
92            yield buf.getvalue()
93    finally:
94        try:
95            fh_.close()
96        except AttributeError:
97            pass
98