diff --git a/README.rst b/README.rst index 993b9ec..74d8800 100644 --- a/README.rst +++ b/README.rst @@ -135,6 +135,10 @@ running ``atrcopy``. +====================+==========+======+=======+==============================+ | gzip | .gz | Yes | No | Read only | +--------------------+----------+------+-------+------------------------------+ +| bzip2 | .bz2 | Yes | No | Read only | ++--------------------+----------+------+-------+------------------------------+ +| lzma | .xz | Yes | No | Read only | ++--------------------+----------+------+-------+------------------------------+ | Disk Communicator | .dcm | No | No | Recognized but unimplemented | +--------------------+----------+------+-------+------------------------------+ diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py index bd7b59f..3c0379e 100644 --- a/atrcopy/__init__.py +++ b/atrcopy/__init__.py @@ -20,7 +20,7 @@ from .kboot import KBootImage, add_xexboot_header from .segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, SegmentedFileSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_style, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments, SegmentList, get_style_mask, get_style_bits from .spartados import SpartaDosDiskImage from .cartridge import A8CartHeader, AtariCartImage -from .parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, guess_container, iter_known_segment_parsers, mime_parse_order, parsers_for_filename +from .parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, guess_container, iter_parsers, iter_known_segment_parsers, mime_parse_order, parsers_for_filename from .magic import guess_detail_for_mime from .utils import to_numpy, text_to_int diff --git a/atrcopy/container.py b/atrcopy/container.py index d357513..499dfe2 100644 --- a/atrcopy/container.py +++ b/atrcopy/container.py @@ -1,4 +1,6 @@ import gzip +import bz2 +import lzma import io import numpy as np @@ -53,3 +55,25 @@ class GZipContainer(DiskImageContainer): except OSError as e: raise errors.InvalidContainer(e) return unpacked + + +class BZipContainer(DiskImageContainer): + def unpack_bytes(self, byte_data): + try: + buf = io.BytesIO(byte_data) + with bz2.BZ2File(buf, mode='rb') as f: + unpacked = f.read() + except OSError as e: + raise errors.InvalidContainer(e) + return unpacked + + +class LZMAContainer(DiskImageContainer): + def unpack_bytes(self, byte_data): + try: + buf = io.BytesIO(byte_data) + with lzma.LZMAFile(buf, mode='rb') as f: + unpacked = f.read() + except lzma.LZMAError as e: + raise errors.InvalidContainer(e) + return unpacked diff --git a/atrcopy/parsers.py b/atrcopy/parsers.py index 69934a9..864f048 100644 --- a/atrcopy/parsers.py +++ b/atrcopy/parsers.py @@ -10,7 +10,7 @@ from .dos33 import Dos33DiskImage, ProdosDiskImage, Dos33BinFile from .standard_delivery import StandardDeliveryImage from . import errors from .magic import guess_detail_for_mime -from .container import GZipContainer +from . import container from .dcm import DCMContainer import logging @@ -160,7 +160,9 @@ class ProdosSegmentParser(SegmentParser): known_containers = [ - GZipContainer, + container.GZipContainer, + container.BZipContainer, + container.LZMAContainer, DCMContainer, ] @@ -207,7 +209,7 @@ def guess_parser_for_system(mime_base, r): def iter_parsers(r): container = guess_container(r.data) if container is not None: - r = container.unpacked + r = SegmentData(container.unpacked) for mime in mime_parse_order: p = guess_parser_for_mime(mime, r) if p is not None: diff --git a/test/test_container.py b/test/test_container.py new file mode 100644 index 0000000..fd87564 --- /dev/null +++ b/test/test_container.py @@ -0,0 +1,41 @@ +from __future__ import print_function +from builtins import object +import numpy as np + +from mock import * + +from atrcopy import SegmentData, iter_parsers +from atrcopy import errors + + +class BaseContainerTest: + base_path = None + expected_mime = "" + + @pytest.mark.parametrize("ext", ['.gz', '.bz2', '.xz', '.dcm']) + def test_container(self, ext): + pathname = self.base_path + ext + try: + sample_data = np.fromfile(pathname, dtype=np.uint8) + except OSError: + pass + else: + rawdata = SegmentData(sample_data.copy()) + mime, parser = iter_parsers(rawdata) + assert mime == self.expected_mime + assert len(parser.image.files) == self.num_files_in_sample + +class TestContainerAtariDosSDImage(BaseContainerTest): + base_path = "../test_data/container_dos_sd_test1.atr" + expected_mime = "application/vnd.atari8bit.atr" + num_files_in_sample = 5 + +class TestContainerAtariDosEDImage(BaseContainerTest): + base_path = "../test_data/container_dos_ed_test1.atr" + expected_mime = "application/vnd.atari8bit.atr" + num_files_in_sample = 5 + +class TestContainerAtariDosDDImage(BaseContainerTest): + base_path = "../test_data/container_dos_dd_test1.atr" + expected_mime = "application/vnd.atari8bit.atr" + num_files_in_sample = 5 diff --git a/test_data/container_dos_dd_test1.atr.bz2 b/test_data/container_dos_dd_test1.atr.bz2 new file mode 100644 index 0000000..97aea66 Binary files /dev/null and b/test_data/container_dos_dd_test1.atr.bz2 differ diff --git a/test_data/container_dos_dd_test1.atr.gz b/test_data/container_dos_dd_test1.atr.gz new file mode 100644 index 0000000..2e4408c Binary files /dev/null and b/test_data/container_dos_dd_test1.atr.gz differ diff --git a/test_data/container_dos_dd_test1.atr.xz b/test_data/container_dos_dd_test1.atr.xz new file mode 100644 index 0000000..6f29a61 Binary files /dev/null and b/test_data/container_dos_dd_test1.atr.xz differ diff --git a/test_data/container_dos_ed_test1.atr.bz2 b/test_data/container_dos_ed_test1.atr.bz2 new file mode 100644 index 0000000..185efd1 Binary files /dev/null and b/test_data/container_dos_ed_test1.atr.bz2 differ diff --git a/test_data/container_dos_ed_test1.atr.gz b/test_data/container_dos_ed_test1.atr.gz new file mode 100644 index 0000000..fd5c792 Binary files /dev/null and b/test_data/container_dos_ed_test1.atr.gz differ diff --git a/test_data/container_dos_ed_test1.atr.xz b/test_data/container_dos_ed_test1.atr.xz new file mode 100644 index 0000000..0e8e47a Binary files /dev/null and b/test_data/container_dos_ed_test1.atr.xz differ diff --git a/test_data/container_dos_sd_test1.atr.bz2 b/test_data/container_dos_sd_test1.atr.bz2 new file mode 100644 index 0000000..6c06a76 Binary files /dev/null and b/test_data/container_dos_sd_test1.atr.bz2 differ diff --git a/test_data/container_dos_sd_test1.atr.gz b/test_data/container_dos_sd_test1.atr.gz new file mode 100644 index 0000000..86f2471 Binary files /dev/null and b/test_data/container_dos_sd_test1.atr.gz differ diff --git a/test_data/container_dos_sd_test1.atr.xz b/test_data/container_dos_sd_test1.atr.xz new file mode 100644 index 0000000..57f45d2 Binary files /dev/null and b/test_data/container_dos_sd_test1.atr.xz differ