From c0340a1807ba5f3e3614d4d9e932b7256bcc9058 Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Tue, 21 Feb 2017 19:49:03 -0800
Subject: [PATCH] Fixed file number in Atari DOS files

---
 atrcopy/ataridos.py   | 6 +++---
 atrcopy/diskimages.py | 6 ++++--
 test/test_add_file.py | 5 ++++-
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index b8e3dab..07617e8 100644
--- a/atrcopy/ataridos.py
+++ b/atrcopy/ataridos.py
@@ -19,9 +19,9 @@ class AtariDosWriteableSector(WriteableSector):
         self._next_sector_num = value
         index = self.sector_size - 3
         hi, lo = divmod(value, 256)
-        self.data[index] = self.used
+        self.data[index + 0] = (self.file_num << 2) | (hi & 0x03)
         self.data[index + 1] = lo
-        self.data[index + 2] = hi
+        self.data[index + 2] = self.used
         log.debug("sector metadata for %d: %s" % (self._sector_num, self.data[index:index + 3]))
         # file number will be added later when known.
 
@@ -177,7 +177,7 @@ class AtariDosDirent(object):
     def process_raw_sector(self, image, raw):
         file_num = raw[-3] >> 2
         if file_num != self.file_num:
-            raise FileNumberMismatchError164()
+            raise FileNumberMismatchError164("Expecting file %d, found %d" % (self.file_num, file_num))
         self.sectors_seen.add(self.current_sector)
         next_sector = ((raw[-3] & 0x3) << 8) + raw[-2]
         if next_sector in self.sectors_seen:
diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py
index be6230e..bdf71f3 100644
--- a/atrcopy/diskimages.py
+++ b/atrcopy/diskimages.py
@@ -387,7 +387,7 @@ class DiskImageBase(object):
         sector_list = self.sector_list_class(self.bytes_per_sector, self.payload_bytes_per_sector, data, self.writeable_sector_class)
         vtoc_segments = self.get_vtoc_segments()
         vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
-        sector_list.calc_sector_map(vtoc)
+        sector_list.calc_sector_map(dirent, vtoc)
         directory.save_dirent(dirent, sector_list)
         self.write_sector_list(sector_list)
         self.write_sector_list(vtoc)
@@ -406,6 +406,7 @@ class WriteableSector(object):
         self._sector_num = -1
         self._next_sector = 0
         self.sector_size = sector_size
+        self.file_num = 0
         self.data = np.zeros([sector_size], dtype=np.uint8)
         self.used = 0
         self.ptr = self.used
@@ -612,7 +613,7 @@ class SectorList(BaseSectorList):
             self.sectors.append(sector)
             index += count
 
-    def calc_sector_map(self, vtoc):
+    def calc_sector_map(self, dirent, vtoc):
         """ Map out the sectors and link the sectors together
 
         raises NotEnoughSpaceOnDisk if the whole file won't fit. It will not
@@ -627,6 +628,7 @@ class SectorList(BaseSectorList):
         last_sector = None
         for sector, sector_num in zip(self.sectors, order):
             sector.sector_num = sector_num
+            sector.file_num = dirent.file_num
             self.file_length += sector.used
             if last_sector is not None:
                 last_sector.next_sector_num = sector_num
diff --git a/test/test_add_file.py b/test/test_add_file.py
index d7bbc45..2ff499d 100644
--- a/test/test_add_file.py
+++ b/test/test_add_file.py
@@ -14,10 +14,13 @@ class TestAtariDosSDImage(object):
     def test_small(self):
         assert len(self.image.files) == 5
 
-        data = [0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2]
+        data = np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8)
         self.image.write_file("TEST.XEX", None, data)
         assert len(self.image.files) == 6
 
+        data2 = self.image.find_file("TEST.XEX")
+        assert data.tostring() == data2
+
 
 if __name__ == "__main__":
     t = TestAtariDosFile()