Reimplement XZ compression (#28)

MANIFEST.in

@@ -17,6 +17,7 @@ include LICENSE
 include README.rst
 
 include get_mozilla_keys.sh
+include update-requirements.sh
 include requirements.txt
 include requirements.in
 include test-requirements.txt

get_mozilla_keys.sh

@@ -11,7 +11,6 @@ function get_key() {
     echo -n "$name = b\"\"\""
     curl -s $url | openssl x509 -inform DER -pubkey -noout | head -c -1
     echo '"""'
-    echo
 }
 
 (
@@ -20,11 +19,16 @@ echo "# Automatically generated - do not edit!"
 echo "#"
 echo "# flake8: noqa"
 get_key "release_primary.der" "release1"
+echo
 get_key "release_secondary.der" "release2"
+echo
 
 get_key "nightly_aurora_level3_primary.der" "nightly1"
+echo
 get_key "nightly_aurora_level3_secondary.der" "nightly2"
+echo
 
 get_key "dep1.der" "dep1"
+echo
 get_key "dep2.der" "dep2"
-) > mardor/mozilla.py
+) > src/mardor/mozilla.py

requirements.in

@@ -1,5 +1,4 @@
 cryptography
 construct
-enum34; python_version < '3.4'
 click
 backports.lzma; python_version < '3.3'

requirements.txt

@@ -5,12 +5,12 @@
 #    pip-compile --output-file requirements.txt requirements.in
 #
 asn1crypto==0.22.0        # via cryptography
-backports.lzma==0.0.6 ; python_version < "3.3"
+backports.lzma==0.0.8 ; python_version < "3.3"
 cffi==1.10.0              # via cryptography
 click==6.7
 construct==2.8.12
-cryptography==2.0
-enum34==1.1.6 ; python_version < "3.4"
+cryptography==2.0.2
+enum34==1.1.6             # via cryptography
 idna==2.5                 # via cryptography
 ipaddress==1.0.18         # via cryptography
 pycparser==2.18           # via cffi

src/mardor/cli.py

@@ -12,7 +12,6 @@
 
 import mardor.mozilla
 from mardor.reader import MarReader
-from mardor.signing import SigningAlgo
 from mardor.signing import get_keysize
 from mardor.writer import MarWriter
 
@@ -48,6 +47,8 @@ def build_argparser():
                         const="bz2", help="compress/decompress members with BZ2")
     parser.add_argument("-J", "--xz", action="store_const", dest="compression",
                         const="xz", help="compress/decompress archive with XZ")
+    parser.add_argument("--auto", action="store_const", dest="compression",
+                        const="auto", help="automatically decompress contents")
 
     parser.add_argument("-k", "--keyfiles", dest="keyfiles", action='append',
                         help="sign/verify with given key(s)")
@@ -106,8 +107,6 @@ def do_list(marfile, detailed=False):
     with open(marfile, 'rb') as f:
         with MarReader(f) as m:
             if detailed:
-                if m.mardata.is_compressed:
-                    yield "MAR data is XZ compressed"
                 if m.mardata.signatures:
                     yield "Signature block found with {} signature".format(m.mardata.signatures.count)
                     for s in m.mardata.signatures.sigs:
@@ -131,12 +130,11 @@ def do_list(marfile, detailed=False):
 def do_create(marfile, files, compress, productversion=None, channel=None,
               signing_key=None, signing_algorithm=None):
     """Create a new MAR file."""
-    xz_compression = compress == mardor.writer.Compression.xz
     with open(marfile, 'w+b') as f:
         with MarWriter(f, productversion=productversion, channel=channel,
                        signing_key=signing_key,
                        signing_algorithm=signing_algorithm,
-                       xz_compression=xz_compression) as m:
+                       ) as m:
             for f in files:
                 m.add(f, compress=compress)
 
@@ -160,17 +158,13 @@ def main(argv=None):
         parser.error("Must specify a key file when verifying")
 
     if args.extract:
-        marfile = os.path.abspath(args.extract)
-        if args.compression == 'bz2':
-            decompress = mardor.reader.Decompression.bz2
-        elif args.compression == 'xz':
-            decompress = mardor.reader.Decompression.xz
-        else:
-            decompress = None
+        if args.compression not in (None, 'bz2', 'xz', 'auto'):
+            parser.error('Unsupported compression type')
 
+        marfile = os.path.abspath(args.extract)
         if args.chdir:
             os.chdir(args.chdir)
-        do_extract(marfile, os.getcwd(), decompress)
+        do_extract(marfile, os.getcwd(), args.compression)
 
     elif args.verify:
         if do_verify(args.verify, args.keyfiles):
@@ -187,21 +181,16 @@ def main(argv=None):
         print("\n".join(do_list(args.list_detailed, detailed=True)))
 
     elif args.create:
-        if args.compression == 'bz2':
-            compress = mardor.writer.Compression.bz2
-        elif args.compression == 'xz':
-            compress = mardor.writer.Compression.xz
-        else:
-            compress = None
-
+        if args.compression not in (None, 'bz2', 'xz'):
+            parser.error('Unsupported compression type')
         marfile = os.path.abspath(args.create)
         if args.keyfiles:
             signing_key = open(args.keyfiles[0], 'rb').read()
             bits = get_keysize(signing_key)
             if bits == 2048:
-                signing_algorithm = SigningAlgo.SHA1
+                signing_algorithm = 'sha1'
             elif bits == 4096:
-                signing_algorithm = SigningAlgo.SHA384
+                signing_algorithm = 'sha384'
             else:
                 parser.error("Unsupported key size {} from key {}".format(bits, args.keyfiles[0]))
 
@@ -212,7 +201,7 @@ def main(argv=None):
 
         if args.chdir:
             os.chdir(args.chdir)
-        do_create(marfile, args.files, compress,
+        do_create(marfile, args.files, args.compression,
                   productversion=args.productversion, channel=args.channel,
                   signing_key=signing_key, signing_algorithm=signing_algorithm)
 

src/mardor/format.py

@@ -111,5 +111,4 @@ def _has_sigs(ctx):
     "data_offset" / Computed(lambda ctx: ctx.index.entries[0].offset),
     "data_length" / Computed(this.header.index_offset - this.data_offset),
     "data_header" / Pointer(this.data_offset, Bytes(6)),
-    "is_compressed" / Computed(this.data_header == b'\xfd7zXZ\x00'),
 )

src/mardor/reader.py

@@ -8,42 +8,24 @@
 """
 
 import os
-import tempfile
-from enum import Enum
 
 from cryptography.exceptions import InvalidSignature
 
 from mardor.format import mar
-from mardor.signing import SigningAlgo
 from mardor.signing import get_signature_data
 from mardor.signing import make_verifier_v1
 from mardor.signing import make_verifier_v2
 from mardor.utils import auto_decompress_stream
 from mardor.utils import bz2_decompress_stream
 from mardor.utils import file_iter
+from mardor.utils import guess_compression
 from mardor.utils import mkdir
 from mardor.utils import safejoin
 from mardor.utils import takeexactly
 from mardor.utils import write_to_file
 from mardor.utils import xz_decompress_stream
 
 
-class Decompression(Enum):
-    """
-    Enum representing different decompression options.
-
-    none: don't decompress
-    auto: automatically decompress depending on specific format
-    bz2: decompress using bz2
-    xz: decompress using xz
-    """
-
-    none = None
-    auto = 1
-    bz2 = 2
-    xz = 3
-
-
 class MarReader(object):
     """Support for reading, extracting, and verifying MAR files.
 
@@ -63,37 +45,9 @@ def __init__(self, fileobj):
                 the MAR data will be read from. This object must also be
                 seekable (i.e.  support .seek() and .tell()).
         """
-        self._raw_fileobj = fileobj
-        self._decompressed_fileobj = None
-
-        self.mardata = mar.parse_stream(self._raw_fileobj)
-        self.is_compressed = self.mardata.is_compressed
-
-    @property
-    def fileobj(self):
-        if not self.is_compressed:
-            return self._raw_fileobj
-
-        if not self._decompressed_fileobj:
-            self._decompressed_fileobj = self.decompress()
-        return self._decompressed_fileobj
-
-    def decompress(self):
-        """Decompress the compressed data section of the mar file
-        Return a fileobject pointing to the decompressed data.
-        """
-        dst = tempfile.TemporaryFile()
-        dst.seek(self.mardata.data_offset)
+        self.fileobj = fileobj
 
-        self._raw_fileobj.seek(self.mardata.data_offset)
-
-        stream = takeexactly(file_iter(self._raw_fileobj), self.mardata.data_length)
-        stream = xz_decompress_stream(stream)
-
-        write_to_file(stream, dst)
-
-        dst.seek(0)
-        return dst
+        self.mardata = mar.parse_stream(self.fileobj)
 
     def __enter__(self):
         """Support the context manager protocol."""
@@ -103,35 +57,75 @@ def __exit__(self, type_, value, tb):
         """Support the context manager protocol."""
         pass
 
-    def extract_entry(self, e, decompress=Decompression.auto):
+    @property
+    def compression_type(self):
+        """Returns the latest compresion type used in this MAR.
+
+        Returns:
+            One of None, 'bz2', or 'xz'
+        """
+        best_compression = None
+        for e in self.mardata.index.entries:
+            self.fileobj.seek(e.offset)
+            magic = self.fileobj.read(10)
+            compression = guess_compression(magic)
+            if compression == 'xz':
+                best_compression = 'xz'
+                break
+            elif compression == 'bz2' and best_compression is None:
+                best_compression = 'bz2'
+        return best_compression
+
+    @property
+    def signature_type(self):
+        """Returns the signature type used in this MAR.
+
+        Returns:
+            One of None, 'sha1', or 'sha384'
+        """
+        if not self.mardata.signatures:
+            return None
+
+        for sig in self.mardata.signatures.sigs:
+            if sig.algorithm_id == 1:
+                return 'sha1'
+            elif sig.algorithm_id == 2:
+                return 'sha384'
+        else:
+            return None
+
+    def extract_entry(self, e, decompress='auto'):
         """Yield blocks of data for this entry from this MAR file.
 
         Args:
             e (:obj:`mardor.format.index_entry`): An index_entry object that
                 refers to this file's size and offset inside the MAR file.
             path (str): Where on disk to extract this file to.
-            decompress (obj, optional): Controls whether files are decompressed
-                when extracted. Must be an instance of Decompression. defaults
-                to Decompression.auto
+            decompress (str, optional): Controls whether files are decompressed
+                when extracted. Must be one of None, 'auto', 'bz2', or 'xz'.
+                Defaults to 'auto'
 
         Yields:
             Blocks of data for `e`
         """
         self.fileobj.seek(e.offset)
         stream = file_iter(self.fileobj)
         stream = takeexactly(stream, e.size)
-        if decompress == Decompression.auto:
+        if decompress == 'auto':
             stream = auto_decompress_stream(stream)
-        elif decompress == Decompression.bz2:
+        elif decompress == 'bz2':
             stream = bz2_decompress_stream(stream)
-        elif decompress == Decompression.xz:
-            # Nothing to do here since we've already decompressed the XZ chunk
+        elif decompress == 'xz':
+            stream = xz_decompress_stream(stream)
+        elif decompress is None:
             pass
+        else:
+            raise ValueError("Unsupported decompression type: {}".format(decompress))
 
         for block in stream:
             yield block
 
-    def extract(self, destdir, decompress=Decompression.auto):
+    def extract(self, destdir, decompress='auto'):
         """Extract the entire MAR file into a directory.
 
         Args:
@@ -166,10 +160,10 @@ def verify(self, verify_key):
 
         verifiers = []
         for sig in self.mardata.signatures.sigs:
-            if sig.algorithm_id == SigningAlgo.SHA1:
+            if sig.algorithm_id == 1:
                 verifier = make_verifier_v1(verify_key, sig.signature)
                 verifiers.append(verifier)
-            elif sig.algorithm_id == SigningAlgo.SHA384:
+            elif sig.algorithm_id == 2:
                 verifier = make_verifier_v2(verify_key, sig.signature)
                 verifiers.append(verifier)
             else:

src/mardor/signing.py

@@ -2,8 +2,6 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 """Signing, verification and key support for MAR files."""
-from enum import IntEnum
-
 from construct import Int32ub
 from construct import Int64ub
 from cryptography.hazmat.backends import default_backend
@@ -16,17 +14,6 @@
 from mardor.utils import file_iter
 
 
-class SigningAlgo(IntEnum):
-    """
-    Enum representing supported signing algorithms.
-
-    SHA1: RSA-PKCS1-SHA1 using 2048 bit key
-    SHA384: RSA-PKCS1-SHA384 using 4096 bit key
-    """
-    SHA1 = 1
-    SHA384 = 2
-
-
 def get_publickey(keydata):
     try:
         key = serialization.load_pem_public_key(