From 68de323d3fff3d2eaf39ff81aa09cb73cad0a443 Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Wed, 1 Jun 2016 15:37:28 -0700
Subject: [PATCH] Added better error checking for atari DOS executables

---
 atrcopy/ataridos.py   |  7 +++++--
 atrcopy/errors.py     |  3 +++
 atrcopy/utils.py      |  2 ++
 test/conftest.py      |  4 ++++
 test/mock.py          | 30 ++++++++++++++++++++++++++++++
 test/test_ataridos.py | 37 +++++++++++++++++++++++++++++++++++++
 6 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 test/conftest.py
 create mode 100644 test/mock.py
 create mode 100644 test/test_ataridos.py

diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index 6d3cb51..6395b9e 100644
--- a/atrcopy/ataridos.py
+++ b/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:
diff --git a/atrcopy/errors.py b/atrcopy/errors.py
index c844c34..90726b0 100644
--- a/atrcopy/errors.py
+++ b/atrcopy/errors.py
@@ -4,6 +4,9 @@ class AtrError(RuntimeError):
 class InvalidAtrHeader(AtrError):
     pass
 
+class InvalidCartHeader(AtrError):
+    pass
+
 class InvalidDiskImage(AtrError):
     pass
 
diff --git a/atrcopy/utils.py b/atrcopy/utils.py
index 6224653..c0432e1 100644
--- a/atrcopy/utils.py
+++ b/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")
 
 
diff --git a/test/conftest.py b/test/conftest.py
new file mode 100644
index 0000000..82a6b2f
--- /dev/null
+++ b/test/conftest.py
@@ -0,0 +1,4 @@
+import pytest
+def pytest_addoption(parser):
+    parser.addoption("--runslow", action="store_true",
+        help="run slow tests")
diff --git a/test/mock.py b/test/mock.py
new file mode 100644
index 0000000..397fc21
--- /dev/null
+++ b/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
diff --git a/test/test_ataridos.py b/test/test_ataridos.py
new file mode 100644
index 0000000..c5ea16c
--- /dev/null
+++ b/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()