From 210b0a74269e4d5d3fa6465565811233068c94bb Mon Sep 17 00:00:00 2001
From: Thomas Daede <daede003@umn.edu>
Date: Mon, 2 Oct 2023 23:13:12 -0700
Subject: [PATCH] D88: Add read support for multiple disks in one file.

---
 src/greaseweazle/image/d88.py     | 33 ++++++++++++++++++++++++-------
 src/greaseweazle/image/image.py   |  3 ++-
 src/greaseweazle/tools/convert.py |  4 +++-
 src/greaseweazle/tools/write.py   |  4 +++-
 4 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/src/greaseweazle/image/d88.py b/src/greaseweazle/image/d88.py
index 3b90401d..0d7e2866 100644
--- a/src/greaseweazle/image/d88.py
+++ b/src/greaseweazle/image/d88.py
@@ -36,9 +36,32 @@ def remove_duplicate_sectors(secs) ->  List[Tuple[int,int,int,int,bytes]]:
         return new_secs
 
     @classmethod
-    def from_file(cls, name: str, _fmt) -> Image:
+    def from_file(cls, name: str, _fmt, index: int = -1) -> Image:
 
         with open(name, "rb") as f:
+            f.seek(0, os.SEEK_END)
+            file_size = f.tell()
+            f.seek(0)
+            disk_offsets: List[int] = []
+            current_disk_index = 0
+            while 1:
+                disk_offsets.append(f.tell())
+                header = struct.unpack('<16sB9xBBL', f.read(32))
+                _, _, _, _, disk_size = header
+                if f.tell() + disk_size >= file_size:
+                    break
+                f.seek(disk_offsets[current_disk_index] + disk_size)
+                current_disk_index += 1
+
+            if index < 0 and len(disk_offsets) > 1:
+                print('D88: Warning: Multiple disks found in image, '
+                      'only using first.')
+                index = 0
+            if index >= len(disk_offsets):
+                raise error.Fatal("D88: No disk with index %d in image file."
+                                  % index)
+
+            f.seek(disk_offsets[index])
             header = struct.unpack('<16sB9xBBL', f.read(32))
             disk_name, terminator, write_prot, media_flag, disk_size = header
             track_table = [x[0] for x in struct.iter_unpack('<L', f.read(640))]
@@ -47,18 +70,14 @@ def from_file(cls, name: str, _fmt) -> Image:
                                     struct.iter_unpack('<L', f.read(16))])
             elif track_table[0] != 672:
                 raise error.Fatal("D88: Unsupported track table length.")
-            f.seek(0, os.SEEK_END)
-            if f.tell() != disk_size:
-                print('D88: Warning: Multiple disks found in image, '
-                      'only using first.')
 
             d88 = cls(name)
 
             for track_index, track_offset in enumerate(track_table):
                 if track_offset == 0:
                     continue
-                f.seek(track_offset)
-                if f.tell() >= disk_size:
+                f.seek(disk_offsets[index] + track_offset)
+                if f.tell() >= disk_offsets[index] + disk_size:
                     continue
                 cyl = track_index // 2
                 head = track_index % 2
diff --git a/src/greaseweazle/image/image.py b/src/greaseweazle/image/image.py
index 31be5e16..ccf76fbc 100644
--- a/src/greaseweazle/image/image.py
+++ b/src/greaseweazle/image/image.py
@@ -67,7 +67,8 @@ def max_cylinder(self):
 
     ## Read support:
     @classmethod
-    def from_file(cls, name: str, fmt: Optional[codec.DiskDef]) -> Image:
+    def from_file(cls, name: str, fmt: Optional[codec.DiskDef],
+                  index: Optional[int]) -> Image:
         raise NotImplementedError
 
     def get_track(self, cyl: int, side: int) -> Optional[HasFlux]:
diff --git a/src/greaseweazle/tools/convert.py b/src/greaseweazle/tools/convert.py
index 5e62edb6..f8cb44aa 100644
--- a/src/greaseweazle/tools/convert.py
+++ b/src/greaseweazle/tools/convert.py
@@ -23,7 +23,7 @@
 plls = track.plls
 
 def open_input_image(args, image_class: Type[Image]) -> Image:
-    return image_class.from_file(args.in_file, args.fmt_cls)
+    return image_class.from_file(args.in_file, args.fmt_cls, int(args.disk_index))
 
 
 def open_output_image(args, image_class: Type[Image]) -> Image:
@@ -117,6 +117,8 @@ def main(argv) -> None:
                                  epilog=epilog)
     parser.add_argument("--diskdefs", help="disk definitions file")
     parser.add_argument("--format", help="disk format")
+    parser.add_argument("--disk-index", help="disk index within input "
+                        "image file", default=-1)
     parser.add_argument("--tracks", type=util.TrackSet,
                         help="which tracks to read & convert from input",
                         metavar="TSPEC")
diff --git a/src/greaseweazle/tools/write.py b/src/greaseweazle/tools/write.py
index 3abf820f..0bc60c29 100644
--- a/src/greaseweazle/tools/write.py
+++ b/src/greaseweazle/tools/write.py
@@ -22,7 +22,7 @@
 
 # Read and parse the image file.
 def open_image(args, image_class: Type[image.Image]) -> image.Image:
-    return image_class.from_file(args.file, args.fmt_cls)
+    return image_class.from_file(args.file, args.fmt_cls, int(args.disk_index))
 
 # write_from_image:
 # Writes the specified image file to floppy disk.
@@ -188,6 +188,8 @@ def main(argv) -> None:
                         help="drive to read")
     parser.add_argument("--diskdefs", help="disk definitions file")
     parser.add_argument("--format", help="disk format")
+    parser.add_argument("--disk-index", help="disk index within image file",
+                        default=-1)
     parser.add_argument("--tracks", type=util.TrackSet, metavar="TSPEC",
                         help="which tracks to write")
     parser.add_argument("--pre-erase", action="store_true",