diff --git a/Tests/images/hopper_naxis_zero.fits b/Tests/images/hopper_naxis_zero.fits new file mode 100644 index 000000000..580cf3a2c Binary files /dev/null and b/Tests/images/hopper_naxis_zero.fits differ diff --git a/Tests/test_file_fitsstub.py b/Tests/test_file_fitsstub.py index 6dc7c4602..c77457947 100644 --- a/Tests/test_file_fitsstub.py +++ b/Tests/test_file_fitsstub.py @@ -1,3 +1,5 @@ +from io import BytesIO + import pytest from PIL import FitsStubImagePlugin, Image @@ -11,10 +13,8 @@ def test_open(): # Assert assert im.format == "FITS" - - # Dummy data from the stub - assert im.mode == "F" - assert im.size == (1, 1) + assert im.size == (128, 128) + assert im.mode == "L" def test_invalid_file(): @@ -35,6 +35,21 @@ def test_load(): im.load() +def test_truncated_fits(): + # No END to headers + image_data = b"SIMPLE = T" + b" " * 50 + b"TRUNCATE" + with pytest.raises(OSError): + FitsStubImagePlugin.FITSStubImageFile(BytesIO(image_data)) + + +def test_naxis_zero(): + # This test image has been manually hexedited + # to set the number of data axes to zero + with pytest.raises(ValueError): + with Image.open("Tests/images/hopper_naxis_zero.fits"): + pass + + def test_save(): # Arrange with Image.open(TEST_FILE) as im: diff --git a/src/PIL/FitsStubImagePlugin.py b/src/PIL/FitsStubImagePlugin.py index c2ce8651c..a3a94cf4b 100644 --- a/src/PIL/FitsStubImagePlugin.py +++ b/src/PIL/FitsStubImagePlugin.py @@ -38,21 +38,45 @@ class FITSStubImageFile(ImageFile.StubImageFile): format_description = "FITS" def _open(self): - offset = self.fp.tell() - if not _accept(self.fp.read(6)): - raise SyntaxError("Not a FITS file") + headers = {} + while True: + header = self.fp.read(80) + if not header: + raise OSError("Truncated FITS file") + keyword = header[:8].strip() + if keyword == b"END": + break + value = header[8:].strip() + if value.startswith(b"="): + value = value[1:].strip() + if not headers and (not _accept(keyword) or value != b"T"): + raise SyntaxError("Not a FITS file") + headers[keyword] = value - # FIXME: add more sanity checks here; mandatory header items - # include SIMPLE, BITPIX, NAXIS, etc. + naxis = int(headers[b"NAXIS"]) + if naxis == 0: + raise ValueError("No image data") + elif naxis == 1: + self._size = 1, int(headers[b"NAXIS1"]) + else: + self._size = int(headers[b"NAXIS1"]), int(headers[b"NAXIS2"]) + + number_of_bits = int(headers[b"BITPIX"]) + if number_of_bits == 8: + self.mode = "L" + elif number_of_bits == 16: + self.mode = "I" + # rawmode = "I;16S" + elif number_of_bits == 32: + self.mode = "I" + elif number_of_bits in (-32, -64): + self.mode = "F" + # rawmode = "F" if number_of_bits == -32 else "F;64F" self.fp.seek(offset) - # make something up - self.mode = "F" - self._size = 1, 1 - loader = self._load() if loader: loader.open(self)