Fixes jbig2 writer to write valid jb2 files

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## [Unreleased]
 
 ### Fixed
+- Fix extraction of jbig2 files, which was producing invalid files ([#652](https://github.com/pdfminer/pdfminer.six/issues/652))
 - Fix issue of TypeError: cannot unpack non-iterable PDFObjRef object, when unpacking the value of 'DW2' ([#529](https://github.com/pdfminer/pdfminer.six/pull/529))
 - `PermissionError` when creating temporary filepaths on windows when running tests ([#469](https://github.com/pdfminer/pdfminer.six/issues/469))
 - Fix `.paint_path` logic for handling single line segments and extracting point-on-curve positions of Beziér path commands ([#530](https://github.com/pdfminer/pdfminer.six/pull/530))

pdfminer/image.py

@@ -92,6 +92,14 @@ def export_image(self, image):
                 fp.write(raw_data)
         elif is_jbig2:
             input_stream = BytesIO()
+            global_streams = self.jbig2_global(image)
+            if len(global_streams) == 1:
+                global_stream, = global_streams
+                input_stream.write(global_stream.get_data().rstrip(b'\n'))
+            elif len(global_streams) > 1:
+                raise ValueError('There should never be more than one '
+                                 'JBIG2Globals associated with a JBIG2 '
+                                 'embedded image')
             input_stream.write(image.stream.get_data())
             input_stream.seek(0)
             reader = JBIG2StreamReader(input_stream)
@@ -137,6 +145,15 @@ def is_jbig2_image(image):
                 break
         return is_jbig2
 
+    @staticmethod
+    def jbig2_global(image):
+        global_streams = []
+        filters = image.stream.get_filters()
+        for filter_name, params in filters:
+            if filter_name in LITERALS_JBIG2_DECODE:
+                global_streams.append(params['JBIG2Globals'].resolve())
+        return global_streams
+
     @staticmethod
     def _get_image_extension(image, width, height, is_jbig2):
         filters = image.stream.get_filters()

pdfminer/jbig2.py

@@ -26,12 +26,11 @@
 # segment types
 SEG_TYPE_IMMEDIATE_GEN_REGION = 38
 SEG_TYPE_END_OF_PAGE = 49
-SEG_TYPE_END_OF_FILE = 50
+SEG_TYPE_END_OF_FILE = 51
 
 # file literals
 FILE_HEADER_ID = b'\x97\x4A\x42\x32\x0D\x0A\x1A\x0A'
 FILE_HEAD_FLAG_SEQUENTIAL = 0b00000001
-FILE_HEAD_FLAG_PAGES_UNKNOWN = 0b00000010
 
 
 def bit_set(bit_pos, value):
@@ -196,8 +195,12 @@ def write_segments(self, segments, fix_last_page=True):
 
     def write_file(self, segments, fix_last_page=True):
         header = FILE_HEADER_ID
-        header_flags = FILE_HEAD_FLAG_SEQUENTIAL | FILE_HEAD_FLAG_PAGES_UNKNOWN
+        header_flags = FILE_HEAD_FLAG_SEQUENTIAL
         header += pack(">B", header_flags)
+        # The embedded JBIG2 files in a PDF always
+        # only have one page
+        number_of_pages = pack(">L", 1)
+        header += number_of_pages
         self.stream.write(header)
         data_len = len(header)
 
@@ -207,7 +210,11 @@ def write_file(self, segments, fix_last_page=True):
         for segment in segments:
             seg_num = segment["number"]
 
-        eof_segment = self.get_eof_segment(seg_num + 1)
+        if fix_last_page:
+            seg_num_offset = 2
+        else:
+            seg_num_offset = 1
+        eof_segment = self.get_eof_segment(seg_num + seg_num_offset)
         data = self.encode_segment(eof_segment)
 
         self.stream.write(data)
@@ -252,7 +259,8 @@ def encode_retention_flags(self, value, segment):
         if ref_count <= 4:
             flags_byte = mask_value(REF_COUNT_SHORT_MASK, ref_count)
             for ref_index, ref_retain in enumerate(retain_segments):
-                flags_byte |= 1 << ref_index
+                if ref_retain:
+                    flags_byte |= 1 << ref_index
             flags.append(flags_byte)
         else:
             bytes_count = math.ceil((ref_count + 1) / 8)

tests/test_tools_pdf2txt.py

@@ -1,6 +1,7 @@
 import os
 from shutil import rmtree
 from tempfile import mkdtemp
+import filecmp
 
 import tools.pdf2txt as pdf2txt
 from helpers import absolute_sample_path
@@ -138,9 +139,21 @@ def test_jbig2_image_export(self):
 
         Feature test for: https://github.com/pdfminer/pdfminer.six/pull/46
         """
-        image_files = self.extract_images(
-            absolute_sample_path('../samples/contrib/pdf-with-jbig2.pdf'))
-        assert image_files[0].endswith('.jb2')
+        input_file = absolute_sample_path(
+            '../samples/contrib/pdf-with-jbig2.pdf')
+        output_dir = mkdtemp()
+        with TemporaryFilePath() as output_file_name:
+            commands = ['-o', output_file_name, '--output-dir',
+                        output_dir, input_file]
+            pdf2txt.main(commands)
+        image_files = os.listdir(output_dir)
+        try:
+            assert image_files[0].endswith('.jb2')
+            assert filecmp.cmp(output_dir + '/' + image_files[0],
+                               absolute_sample_path(
+                                   '../samples/contrib/XIPLAYER0.jb2'))
+        finally:
+            rmtree(output_dir)
 
     def test_contrib_matplotlib(self):
         """Test a pdf with Type3 font"""