Added better error checking for atari DOS executables

atrcopy/ataridos.py

@@ -166,18 +166,21 @@ class AtariDosFile(object):
             else:
                 self.segments.append(ObjSegment(r[pos:pos + 1], pos, pos + 1, 0, 1, "Incomplete Data"))
                 break
-            log.debug("header parsing: header=0x%x" % header)
             if header == 0xffff:
                 # Apparently 0xffff header can appear in any segment, not just
                 # the first.  Regardless, it is ignored everywhere.
                 pos += 2
+                first = False
+                continue
             elif first:
                 raise InvalidBinaryFile
-            first = False
+            log.debug("header parsing: header=0x%x" % header)
             if len(b[pos:pos + 4]) < 4:
                 self.segments.append(ObjSegment(r[pos:pos + 4], 0, 0, 0, len(b[pos:pos + 4]), "Short Segment Header"))
                 break
             start, end = b[pos:pos + 4].view(dtype='<u2')
+            if end < start:
+                raise InvalidBinaryFile
             count = end - start + 1
             found = len(b[pos + 4:pos + 4 + count])
             if found < count:

atrcopy/errors.py

@@ -4,6 +4,9 @@ class AtrError(RuntimeError):
 class InvalidAtrHeader(AtrError):
     pass
 
+class InvalidCartHeader(AtrError):
+    pass
+
 class InvalidDiskImage(AtrError):
     pass
 

atrcopy/utils.py

@@ -8,6 +8,8 @@ def to_numpy(value):
         return value
     elif type(value) is types.StringType:
         return np.fromstring(value, dtype=np.uint8)
+    elif type(value) is types.ListType:
+    	return np.asarray(value, dtype=np.uint8)
     raise TypeError("Can't convert to numpy data")
 
 

test/conftest.py

@@ -0,0 +1,4 @@
+import pytest
+def pytest_addoption(parser):
+    parser.addoption("--runslow", action="store_true",
+        help="run slow tests")

test/mock.py

@@ -0,0 +1,30 @@
+from __future__ import print_function
+
+import os
+
+# Include maproom directory so that maproom modules can be imported normally
+import sys
+module_dir = os.path.realpath(os.path.abspath(".."))
+if module_dir not in sys.path:
+    sys.path.insert(0, module_dir)
+
+import pytest
+try:
+    slow = pytest.mark.skipif(
+        not pytest.config.getoption("--runslow"),
+        reason="need --runslow option to run"
+        )
+except AttributeError:
+    # pytest doesn't load the config module when not run using py.test
+    # skip this check when running a test_*.py from the command line
+    import functools
+    slow = lambda a: functools.partial(print, "skipping slow test %s" % repr(a))
+
+# Turn logging on by default at the DEBUG level for tests
+import logging
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+
+import numpy as np
+from numpy.testing import assert_almost_equal

test/test_ataridos.py

@@ -0,0 +1,37 @@
+from mock import *
+
+from atrcopy import SegmentData, AtariDosFile, InvalidBinaryFile
+
+
+class TestAtariDosFile(object):
+    def setup(self):
+        pass
+
+    def test_segment(self):
+        bytes = [0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2]
+        rawdata = SegmentData(bytes)
+        image = AtariDosFile(rawdata)
+        image.parse_segments()
+        assert len(image.segments) == 1
+        assert len(image.segments[0]) == 2
+
+    def test_short_segment(self):
+        bytes = [0xff, 0xff, 0x00, 0x60, 0xff, 0x60, 1, 2]
+        rawdata = SegmentData(bytes)
+        image = AtariDosFile(rawdata)
+        image.parse_segments()
+        assert len(image.segments) == 1
+        assert len(image.segments[0]) == 2
+
+    def test_err_segment(self):
+        bytes = [0xff, 0xff, 0x00, 0x60, 0x00, 0x00, 1, 2]
+        rawdata = SegmentData(bytes)
+        image = AtariDosFile(rawdata)
+        with pytest.raises(InvalidBinaryFile):
+            image.parse_segments()
+
+
+if __name__ == "__main__":
+    t = TestAtariDosFile()
+    t.setup()
+    t.test_segment()