Merge pull request #65 from jcristau/bcj-x86

Don't use x86 lzma filter unconditionally

src/mardor/cli.py

@@ -31,6 +31,13 @@ def build_argparser():
                               help="channel this MAR file is applicable to")
     create_group.add_argument("files", nargs=REMAINDER,
                               help="files to add to the MAR file")
+    create_group.add_argument(
+        "--x86",
+        action="store_const",
+        dest="bcj",
+        const="x86",
+        help="use x86 BCJ filter for XZ compression",
+    )
 
     extract_group = parser.add_argument_group("Extract a MAR file")
     extract_group.add_argument("-x", "--extract", help="extract MAR", metavar="MARFILE")
@@ -186,7 +193,7 @@ def do_list(marfile, detailed=False):
 
 
 def do_create(marfile, files, compress, productversion=None, channel=None,
-              signing_key=None, signing_algorithm=None):
+              signing_key=None, signing_algorithm=None, bcj=None):
     """Create a new MAR file."""
     with open(marfile, 'w+b') as f:
         with MarWriter(f, productversion=productversion, channel=channel,
@@ -194,7 +201,7 @@ def do_create(marfile, files, compress, productversion=None, channel=None,
                        signing_algorithm=signing_algorithm,
                        ) as m:
             for f in files:
-                m.add(f, compress=compress)
+                m.add(f, compress=compress, bcj=bcj)
 
 
 def do_hash(hash_algo, marfile, asn1=False):
@@ -254,6 +261,9 @@ def check_args(parser, args):
     if args.hash and len(args.files) != 1:
         parser.error("Must specify a file to output the hash for")
 
+    if args.create and args.bcj and args.compression != "xz":
+        parser.error("BCJ filter is only valid for XZ compression")
+
 
 def get_key_from_cmdline(parser, args):
     """Return the signing key and signing algoritm from the commandline."""
@@ -309,7 +319,8 @@ def main(argv=None):
             os.chdir(args.chdir)
         do_create(marfile, args.files, args.compression,
                   productversion=args.productversion, channel=args.channel,
-                  signing_key=signing_key, signing_algorithm=signing_algorithm)
+                  signing_key=signing_key, signing_algorithm=signing_algorithm,
+                  bcj=args.bcj)
 
     elif args.hash:
         do_hash(args.hash, args.files[0], args.asn1)

src/mardor/utils.py

@@ -142,7 +142,7 @@ def bz2_decompress_stream(src):
             yield decoded
 
 
-def xz_compress_stream(src):
+def xz_compress_stream(src, bcj=None):
     """Compress data from `src`.
 
     Args:
@@ -152,13 +152,17 @@ def xz_compress_stream(src):
         blocks of compressed data
 
     """
-    compressor = lzma.LZMACompressor(
-        check=lzma.CHECK_CRC64,
-        filters=[
-            {"id": lzma.FILTER_X86},
+    filters = []
+    if bcj == "x86":
+        filters.append({"id": lzma.FILTER_X86})
+    filters.extend([
             {"id": lzma.FILTER_LZMA2,
              "preset": lzma.PRESET_DEFAULT},
         ])
+    compressor = lzma.LZMACompressor(
+        check=lzma.CHECK_CRC64,
+        filters=filters,
+    )
     for block in src:
         encoded = compressor.compress(block)
         if encoded:

src/mardor/writer.py

@@ -114,7 +114,7 @@ class MarWriter(object):
         self.finish()
         self.flush()
 
-    def add(self, path, compress=None):
+    def add(self, path, compress=None, bcj=None):
         """Add `path` to the MAR file.
 
         If `path` is a file, it will be added directly.
@@ -125,45 +125,49 @@ class MarWriter(object):
             path (str): path to file or directory on disk to add to this MAR
                 file
             compress (str): One of 'xz', 'bz2', or None. Defaults to None.
+            bcj (str): If compress is 'xz', one of 'x86' or None.
         """
         if os.path.isdir(path):
-            self.add_dir(path, compress)
+            self.add_dir(path, compress, bcj)
         else:
-            self.add_file(path, compress)
+            self.add_file(path, compress, bcj)
 
-    def add_dir(self, path, compress):
+    def add_dir(self, path, compress, bcj=None):
         """Add all files under directory `path` to the MAR file.
 
         Args:
             path (str): path to directory to add to this MAR file
-            compress (str): One of 'xz', 'bz2', or None. Defaults to None.
+            compress (str): One of 'xz', 'bz2', or None.
+            bcj (str): If compress is 'xz', one of 'x86' or None.
         """
         if not os.path.isdir(path):
             raise ValueError('{} is not a directory'.format(path))
         for root, dirs, files in os.walk(path):
             for f in files:
-                self.add_file(os.path.join(root, f), compress)
+                self.add_file(os.path.join(root, f), compress, bcj)
 
-    def add_fileobj(self, fileobj, path, compress, flags=None):
+    def add_fileobj(self, fileobj, path, compress, flags=None, bcj=None):
         """Add the contents of a file object to the MAR file.
 
         Args:
             fileobj (file-like object): open file object
             path (str): name of this file in the MAR file
-            compress (str): One of 'xz', 'bz2', or None. Defaults to None.
+            compress (str): One of 'xz', 'bz2', or None.
+            bcj (str): If compress is 'xz', one of 'x86' or None.
             flags (int): permission of this file in the MAR file. Defaults to the permissions of `path`
         """
         f = file_iter(fileobj)
         flags = flags or os.stat(path) & 0o777
-        return self.add_stream(f, path, compress, flags)
+        return self.add_stream(f, path, compress, flags, bcj)
 
-    def add_stream(self, stream, path, compress, flags):
+    def add_stream(self, stream, path, compress, flags, bcj=None):
         """Add the contents of an iterable to the MAR file.
 
         Args:
             stream (iterable): yields blocks of data
             path (str): name of this file in the MAR file
             compress (str): One of 'xz', 'bz2', or None. Defaults to None.
+            bcj (str): If compress is 'xz', one of 'x86' or None.
             flags (int): permission of this file in the MAR file
         """
         self.data_fileobj.seek(self.last_offset)
@@ -171,7 +175,7 @@ class MarWriter(object):
         if compress == 'bz2':
             stream = bz2_compress_stream(stream)
         elif compress == 'xz':
-            stream = xz_compress_stream(stream)
+            stream = xz_compress_stream(stream, bcj)
         elif compress is None:
             pass
         else:
@@ -193,12 +197,13 @@ class MarWriter(object):
         self.entries.append(e)
         self.last_offset += e['size']
 
-    def add_file(self, path, compress):
+    def add_file(self, path, compress, bcj=None):
         """Add a single file to the MAR file.
 
         Args:
             path (str): path to a file to add to this MAR file.
-            compress (str): One of 'xz', 'bz2', or None. Defaults to None.
+            compress (str): One of 'xz', 'bz2', or None.
+            bcj (str): If compress is 'xz', one of 'x86' or None.
         """
         if not os.path.isfile(path):
             raise ValueError('{} is not a file'.format(path))
@@ -206,7 +211,7 @@ class MarWriter(object):
 
         with open(path, 'rb') as f:
             flags = os.stat(path).st_mode & 0o777
-            self.add_fileobj(f, path, compress, flags)
+            self.add_fileobj(f, path, compress, flags, bcj)
 
     def write_header(self):
         """Write the MAR header to the file.

tests/test_writer.py

@@ -207,11 +207,14 @@ def test_adddir_as_file(tmpdir):
 def test_xz_writer(tmpdir):
     message_p = tmpdir.join('message.txt')
     message_p.write('hello world')
+    x86_message_p = tmpdir.join('message_x86.txt')
+    x86_message_p.write('hello world')
     mar_p = tmpdir.join('test.mar')
     with mar_p.open('wb') as f:
         with MarWriter(f) as m:
             with tmpdir.as_cwd():
                 m.add('message.txt', compress='xz')
+                m.add('message_x86.txt', compress='xz', bcj='x86')
 
     assert mar_p.size() > 0
 
@@ -219,11 +222,14 @@ def test_xz_writer(tmpdir):
         with MarReader(f) as m:
             assert m.mardata.additional is None
             assert m.mardata.signatures is None
-            assert len(m.mardata.index.entries) == 1
+            assert len(m.mardata.index.entries) == 2
             assert m.mardata.index.entries[0].name == 'message.txt'
+            assert m.mardata.index.entries[1].name == 'message_x86.txt'
             m.extract(str(tmpdir.join('extracted')))
             assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
                     b'hello world')
+            assert (tmpdir.join('extracted', 'message_x86.txt').read('rb') ==
+                    b'hello world')
 
 
 def test_writer_badmode(tmpdir, test_keys):