From f1038329ee6fa98d5156c65cc8c9bfe4ae21b8d8 Mon Sep 17 00:00:00 2001 From: Rob McMullen Date: Wed, 1 Jun 2016 14:54:52 -0700 Subject: [PATCH] Added parsers and MIME type definitions from Omnivore * restructured main script to use parsers, replacing the old logic --- atrcopy/__init__.py | 73 ++++++++++++--------------------- atrcopy/ataridos.py | 2 +- atrcopy/diskimages.py | 18 ++++----- atrcopy/errors.py | 3 ++ atrcopy/kboot.py | 2 +- atrcopy/parsers.py | 93 +++++++++++++++++++++++++++++++++++++++++++ atrcopy/spartados.py | 4 +- 7 files changed, 134 insertions(+), 61 deletions(-) create mode 100644 atrcopy/parsers.py diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py index 92536dd..ac254a8 100644 --- a/atrcopy/__init__.py +++ b/atrcopy/__init__.py @@ -13,6 +13,7 @@ from diskimages import AtrHeader, BootDiskImage, add_atr_header from kboot import KBootImage, add_xexboot_header from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_bit_mask, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments from spartados import SpartaDosDiskImage +from parsers import SegmentParser, DefaultSegmentParser, guess_parser_for, known_segment_parsers, mime_parse_order from utils import to_numpy @@ -69,53 +70,29 @@ def run(): for filename in options.files: with open(filename, "rb") as fh: + if options.verbose: + print "Loading file %s" % filename rawdata = SegmentData(fh.read()) - data = rawdata.get_data() - image = None - if options.debug: - header = AtrHeader(data[0:16]) - image = SpartaDosDiskImage(rawdata, filename) - else: - try: - data = to_numpy(data) - try: - header = AtrHeader(data[0:16]) - for format in [KBootImage, SpartaDosDiskImage, AtariDosDiskImage]: - if options.verbose: print "trying", format.__name__ - try: - image = format(rawdata, filename) - print "%s: %s" % (filename, image) - break - except InvalidDiskImage: - pass - except AtrError: - for format in [AtariDosDiskImage]: - try: - image = format(rawdata, filename) - print "%s: %s" % (filename, image) - break - except: - raise - #pass - except AtrError: - if options.verbose: print "%s: Doesn't look like a supported disk image" % filename - try: - image = AtariDosFile(rawdata) - print "%s:\n%s" % (filename, image) - except InvalidBinaryFile: - if options.verbose: print "%s: Doesn't look like an XEX either" % filename + parser = None + for mime in mime_parse_order: + if options.verbose: + print "Trying MIME type %s" % mime + parser = guess_parser_for(mime, rawdata) + if parser is None: continue - if image is None: - image = BootDiskImage(rawdata, filename) - if options.segments: - image.parse_segments() - print "\n".join([str(a) for a in image.segments]) - elif image.files or options.force: - for dirent in image.files: - try: - process(image, dirent, options) - except FileNumberMismatchError164: - print "Error 164: %s" % str(dirent) - except ByteNotInFile166: - print "Invalid sector for: %s" % str(dirent) - + if options.verbose: + print "Found parser %s" % parser.menu_name + print "%s: %s" % (filename, parser.image) + if options.segments: + print "\n".join([str(a) for a in parser.segments]) + elif parser.image.files or options.force: + for dirent in parser.image.files: + try: + process(parser.image, dirent, options) + except FileNumberMismatchError164: + print "Error 164: %s" % str(dirent) + except ByteNotInFile166: + print "Invalid sector for: %s" % str(dirent) + break + if parser is None: + print "%s: Unknown file type" % filename diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py index 7603361..6d3cb51 100644 --- a/atrcopy/ataridos.py +++ b/atrcopy/ataridos.py @@ -149,7 +149,7 @@ class AtariDosFile(object): self.rawdata = rawdata self.size = len(rawdata) self.segments = [] - self.parse_segments() + self.files = [] def __str__(self): return "\n".join(str(s) for s in self.segments) + "\n" diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py index 8b37ab2..73c818b 100644 --- a/atrcopy/diskimages.py +++ b/atrcopy/diskimages.py @@ -23,12 +23,12 @@ class AtrHeader(object): self.crc = 0 self.unused = 0 self.flags = 0 - self.atr_header_offset = 0 + self.header_offset = 0 self.initial_sector_size = sector_size self.num_initial_sectors = initial_sectors self.max_sectors = 0 if create: - self.atr_header_offset = 16 + self.header_offset = 16 self.check_size(0) if bytes is None: return @@ -42,7 +42,7 @@ class AtrHeader(object): self.crc = int(values[4]) self.unused = int(values[5]) self.flags = int(values[6]) - self.atr_header_offset = 16 + self.header_offset = 16 else: raise InvalidAtrHeader @@ -50,7 +50,7 @@ class AtrHeader(object): return "%s Disk Image (size=%d (%dx%db), crc=%d flags=%d unused=%d)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size, self.crc, self.flags, self.unused) def __len__(self): - return self.atr_header_offset + return self.header_offset def to_array(self): raw = np.zeros([16], dtype=np.uint8) @@ -100,7 +100,7 @@ class AtrHeader(object): else: pos = self.num_initial_sectors * self.initial_sector_size + (sector - 1 - self.num_initial_sectors) * self.sector_size size = self.sector_size - pos += self.atr_header_offset + pos += self.header_offset return pos, size @@ -149,7 +149,7 @@ class DiskImageBase(object): def setup(self): self.size = np.alen(self.bytes) - self.read_atr_header() + self.read_header() self.header.check_size(self.size - len(self.header)) self.check_size() self.get_boot_sector_info() @@ -195,7 +195,7 @@ class DiskImageBase(object): if not self.all_sane: raise InvalidDiskImage("Invalid directory entries; may be boot disk") - def read_atr_header(self): + def read_header(self): bytes = self.bytes[0:16] try: self.header = AtrHeader(bytes) @@ -256,7 +256,7 @@ class DiskImageBase(object): def parse_segments(self): r = self.rawdata - i = self.header.atr_header_offset + i = self.header.header_offset if self.header.image_size > 0: self.segments.append(ObjSegment(r[0:i], 0, 0, 0, i, name="%s Header" % self.header.file_format)) self.segments.append(RawSectorsSegment(r[i:], 1, self.header.max_sectors, self.header.image_size, 128, 3, self.header.sector_size, name="Raw disk sectors")) @@ -327,7 +327,7 @@ class BootDiskImage(DiskImageBase): return start, size = self.header.get_pos(1) b = self.bytes - i = self.header.atr_header_offset + i = self.header.header_offset flag = b[i:i + 2].view(dtype=' image.size or start + count < image.size - 128: self.is_sane = False diff --git a/atrcopy/parsers.py b/atrcopy/parsers.py new file mode 100644 index 0000000..b5dc34f --- /dev/null +++ b/atrcopy/parsers.py @@ -0,0 +1,93 @@ +import numpy as np + +from segments import SegmentData, DefaultSegment +from diskimages import BootDiskImage +from kboot import KBootImage +from ataridos import AtariDosDiskImage, AtariDosFile +from spartados import SpartaDosDiskImage +from errors import * + + +class SegmentParser(object): + menu_name = "" + image_type = None + + def __init__(self, segment_data): + self.image = None + self.segments = [] + self.parse(segment_data) + + def parse(self, r): + self.segments.append(DefaultSegment(r, 0)) + try: + self.image = self.image_type(r) + self.image.parse_segments() + except AtrError: + raise InvalidSegmentParser + self.segments.extend(self.image.segments) + + +class DefaultSegmentParser(SegmentParser): + menu_name = "Raw Data" + + def parse(self, r): + self.segments = [DefaultSegment(r, 0)] + + +class KBootSegmentParser(SegmentParser): + menu_name = "KBoot Disk Image" + image_type = KBootImage + + +class AtariDosSegmentParser(SegmentParser): + menu_name = "Atari DOS Disk Image" + image_type = AtariDosDiskImage + + +class SpartaDosSegmentParser(SegmentParser): + menu_name = "Sparta DOS Disk Image" + image_type = SpartaDosDiskImage + + +class AtariBootDiskSegmentParser(SegmentParser): + menu_name = "Atari Boot Disk Image" + image_type = BootDiskImage + + +class XexSegmentParser(SegmentParser): + menu_name = "XEX (Atari 8-bit executable)" + image_type = AtariDosFile + + +def guess_parser_for(mime, r): + parsers = mime_parsers[mime] + found = None + for parser in parsers: + try: + found = parser(r) + break + except InvalidSegmentParser: + pass + return found + + +mime_parsers = { + "application/vnd.atari8bit.atr": [ + KBootSegmentParser, + SpartaDosSegmentParser, + AtariDosSegmentParser, + AtariBootDiskSegmentParser, + ], + "application/vnd.atari8bit.xex": [ + XexSegmentParser, + ], + } + +mime_parse_order = [ + "application/vnd.atari8bit.atr", + "application/vnd.atari8bit.xex", + ] + +known_segment_parsers = [DefaultSegmentParser] +for mime in mime_parse_order: + known_segment_parsers.extend(mime_parsers[mime]) diff --git a/atrcopy/spartados.py b/atrcopy/spartados.py index dee35ec..509faa5 100644 --- a/atrcopy/spartados.py +++ b/atrcopy/spartados.py @@ -108,13 +108,13 @@ class SpartaDosDirent(AtariDosDirent): class SpartaDosDiskImage(DiskImageBase): - def __init__(self, bytes, style=None): + def __init__(self, *args, **kwargs): self.first_bitmap = 0 self.num_bitmap = 0 self.root_dir = 0 self.root_dir_dirent = None self.fs_version = 0 - DiskImageBase.__init__(self, bytes, style) + DiskImageBase.__init__(self, *args, **kwargs) def __str__(self): return "%s Sparta DOS Format: %d usable sectors (%d free), %d files" % (self.header, self.total_sectors, self.unused_sectors, len(self.files))