Added SegmentData class and updated to version 2.5.0 in preparation for adding comments

This commit is contained in:
Rob McMullen 2016-03-25 16:07:16 -07:00
parent 3c2f6df828
commit c8ccd73779
5 changed files with 116 additions and 80 deletions

View File

@ -1,4 +1,4 @@
__version__ = "2.4.0" __version__ = "2.5.0"
try: try:
import numpy as np import numpy as np
@ -9,7 +9,7 @@ from errors import *
from ataridos import AtariDosDiskImage, AtariDosFile from ataridos import AtariDosDiskImage, AtariDosFile
from diskimages import AtrHeader, BootDiskImage from diskimages import AtrHeader, BootDiskImage
from kboot import KBootImage from kboot import KBootImage
from segments import SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment
from spartados import SpartaDosDiskImage from spartados import SpartaDosDiskImage
from utils import to_numpy from utils import to_numpy
@ -60,12 +60,12 @@ def run():
for filename in options.files: for filename in options.files:
with open(filename, "rb") as fh: with open(filename, "rb") as fh:
data = fh.read() rawdata = SegmentData(fh.read())
data = rawdata.get_data()
image = None image = None
if options.debug: if options.debug:
data = to_numpy(data)
header = AtrHeader(data[0:16]) header = AtrHeader(data[0:16])
image = SpartaDosDiskImage(data, filename) image = SpartaDosDiskImage(rawdata, filename)
else: else:
try: try:
data = to_numpy(data) data = to_numpy(data)
@ -74,7 +74,7 @@ def run():
for format in [KBootImage, SpartaDosDiskImage, AtariDosDiskImage]: for format in [KBootImage, SpartaDosDiskImage, AtariDosDiskImage]:
if options.verbose: print "trying", format.__name__ if options.verbose: print "trying", format.__name__
try: try:
image = format(data, filename) image = format(rawdata, filename)
print "%s: %s" % (filename, image) print "%s: %s" % (filename, image)
break break
except InvalidDiskImage: except InvalidDiskImage:
@ -82,7 +82,7 @@ def run():
except AtrError: except AtrError:
for format in [AtariDosDiskImage]: for format in [AtariDosDiskImage]:
try: try:
image = format(data) image = format(rawdata, filename)
print "%s: %s" % (filename, image) print "%s: %s" % (filename, image)
break break
except: except:
@ -91,13 +91,13 @@ def run():
except AtrError: except AtrError:
if options.verbose: print "%s: Doesn't look like a supported disk image" % filename if options.verbose: print "%s: Doesn't look like a supported disk image" % filename
try: try:
image = AtariDosFile(data) image = AtariDosFile(rawdata)
print "%s:\n%s" % (filename, image) print "%s:\n%s" % (filename, image)
except InvalidBinaryFile: except InvalidBinaryFile:
if options.verbose: print "%s: Doesn't look like an XEX either" % filename if options.verbose: print "%s: Doesn't look like an XEX either" % filename
continue continue
if image is None: if image is None:
image = BootDiskImage(data, filename) image = BootDiskImage(rawdata, filename)
if options.segments: if options.segments:
image.parse_segments() image.parse_segments()
print "\n".join([str(a) for a in image.segments]) print "\n".join([str(a) for a in image.segments])

View File

@ -142,13 +142,11 @@ class AtariDosFile(object):
Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html Ref: http://www.atarimax.com/jindroush.atari.org/afmtexe.html
""" """
def __init__(self, data, style=None): def __init__(self, rawdata):
self.bytes = to_numpy(data) self.rawdata = rawdata
self.size = np.alen(self.bytes) self.bytes = rawdata.get_data()
if style is None: self.style = rawdata.get_style()
self.style = np.zeros(self.size, dtype=np.uint8) self.size = len(rawdata)
else:
self.style = style
self.segments = [] self.segments = []
self.parse_segments() self.parse_segments()
@ -156,6 +154,7 @@ class AtariDosFile(object):
return "\n".join(str(s) for s in self.segments) + "\n" return "\n".join(str(s) for s in self.segments) + "\n"
def parse_segments(self): def parse_segments(self):
r = self.rawdata
b = self.bytes b = self.bytes
s = self.style s = self.style
pos = 0 pos = 0
@ -164,7 +163,7 @@ class AtariDosFile(object):
if pos + 1 < self.size: if pos + 1 < self.size:
header, = b[pos:pos+2].view(dtype='<u2') header, = b[pos:pos+2].view(dtype='<u2')
else: else:
self.segments.append(ObjSegment(b[pos:pos + 1], s[pos:pos + 1], pos, pos + 1, 0, 1, "Incomplete Data")) self.segments.append(ObjSegment(r[pos:pos + 1], pos, pos + 1, 0, 1, "Incomplete Data"))
break break
if header == 0xffff: if header == 0xffff:
# Apparently 0xffff header can appear in any segment, not just # Apparently 0xffff header can appear in any segment, not just
@ -174,25 +173,25 @@ class AtariDosFile(object):
raise InvalidBinaryFile raise InvalidBinaryFile
first = False first = False
if len(b[pos:pos + 4]) < 4: if len(b[pos:pos + 4]) < 4:
self.segments.append(ObjSegment(b[pos:pos + 4], s[pos:pos + 4], 0, 0, "Short Segment Header")) self.segments.append(ObjSegment(r[pos:pos + 4], 0, 0, "Short Segment Header"))
break break
start, end = b[pos:pos + 4].view(dtype='<u2') start, end = b[pos:pos + 4].view(dtype='<u2')
count = end - start + 1 count = end - start + 1
found = len(b[pos + 4:pos + 4 + count]) found = len(b[pos + 4:pos + 4 + count])
if found < count: if found < count:
self.segments.append(ObjSegment(b[pos + 4:pos + 4 + count], s[pos + 4:pos + 4 + count], pos, pos + 4, start, end, "Incomplete Data")) self.segments.append(ObjSegment(r[pos + 4:pos + 4 + count], pos, pos + 4, start, end, "Incomplete Data"))
break break
self.segments.append(ObjSegment(b[pos + 4:pos + 4 + count], s[pos + 4:pos + 4 + count], pos, pos + 4, start, end)) self.segments.append(ObjSegment(r[pos + 4:pos + 4 + count], pos, pos + 4, start, end))
pos += 4 + count pos += 4 + count
class AtariDosDiskImage(DiskImageBase): class AtariDosDiskImage(DiskImageBase):
def __init__(self, bytes, style=None): def __init__(self, *args, **kwargs):
self.first_vtoc = 360 self.first_vtoc = 360
self.num_vtoc = 1 self.num_vtoc = 1
self.vtoc2 = 0 self.vtoc2 = 0
self.first_data_after_vtoc = 369 self.first_data_after_vtoc = 369
DiskImageBase.__init__(self, bytes, style) DiskImageBase.__init__(self, *args, **kwargs)
def __str__(self): def __str__(self):
return "%s Atari DOS Format: %d usable sectors (%d free), %d files" % (self.header, self.total_sectors, self.unused_sectors, len(self.files)) return "%s Atari DOS Format: %d usable sectors (%d free), %d files" % (self.header, self.total_sectors, self.unused_sectors, len(self.files))
@ -266,34 +265,33 @@ class AtariDosDiskImage(DiskImageBase):
if flag == 0: if flag == 0:
num = int(values[1]) num = int(values[1])
addr = int(values[2]) addr = int(values[2])
bytes, style = self.get_sectors(1, num) s = self.get_sector_slice(1, num)
header = ObjSegment(bytes[0:20], style[0:20], 0, 0, addr, addr + 20, name="Boot Header") r = self.rawdata[s]
sectors = ObjSegment(bytes, style, 0, 0, addr, addr + len(bytes), name="Boot Sectors") header = ObjSegment(r[0:20], 0, 0, addr, addr + 20, name="Boot Header")
code = ObjSegment(bytes[20:], style[20:], 0, 0, addr + 20, addr + len(bytes), name="Boot Code") sectors = ObjSegment(r, 0, 0, addr, addr + len(r), name="Boot Sectors")
code = ObjSegment(r[20:], 0, 0, addr + 20, addr + len(r), name="Boot Code")
segments = [sectors, header, code] segments = [sectors, header, code]
return segments return segments
def get_vtoc_segments(self): def get_vtoc_segments(self):
b = self.bytes r = self.rawdata
s = self.style
segments = [] segments = []
addr = 0 addr = 0
start, count = self.get_contiguous_sectors(self.first_vtoc, self.num_vtoc) start, count = self.get_contiguous_sectors(self.first_vtoc, self.num_vtoc)
segment = RawSectorsSegment(b[start:start+count], s[start:start+count], self.first_vtoc, self.num_vtoc, count, 128, 3, self.header.sector_size, name="VTOC") segment = RawSectorsSegment(r[start:start+count], self.first_vtoc, self.num_vtoc, count, 128, 3, self.header.sector_size, name="VTOC")
segments.append(segment) segments.append(segment)
if self.vtoc2 > 0: if self.vtoc2 > 0:
start, count = self.get_contiguous_sectors(self.vtoc2, 1) start, count = self.get_contiguous_sectors(self.vtoc2, 1)
segment = RawSectorsSegment(b[start:start+count], s[start:start+count], self.vtoc2, 1, count, self.header.sector_size, name="VTOC2") segment = RawSectorsSegment(r[start:start+count], self.vtoc2, 1, count, self.header.sector_size, name="VTOC2")
segments.append(segment) segments.append(segment)
return segments return segments
def get_directory_segments(self): def get_directory_segments(self):
b = self.bytes r = self.rawdata
s = self.style
segments = [] segments = []
addr = 0 addr = 0
start, count = self.get_contiguous_sectors(361, 8) start, count = self.get_contiguous_sectors(361, 8)
segment = RawSectorsSegment(b[start:start+count], s[start:start+count], 361, 8, count, 128, 3, self.header.sector_size, name="Directory") segment = RawSectorsSegment(r[start:start+count], 361, 8, count, 128, 3, self.header.sector_size, name="Directory")
segments.append(segment) segments.append(segment)
return segments return segments
@ -309,7 +307,7 @@ class AtariDosDiskImage(DiskImageBase):
name = "%s %ds@%d" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector) name = "%s %ds@%d" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector)
verbose_name = "%s (%d sectors, first@%d) %s" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector, dirent.verbose_info) verbose_name = "%s (%d sectors, first@%d) %s" % (dirent.get_filename(), dirent.num_sectors, dirent.starting_sector, dirent.verbose_info)
print verbose_name print verbose_name
segment = IndexedByteSegment(self.bytes, self.style, byte_order, name=name, verbose_name=verbose_name) segment = IndexedByteSegment(self.rawdata, byte_order, name=name, verbose_name=verbose_name)
else: else:
segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename()) segment = EmptySegment(self.rawdata, name=dirent.get_filename())
return segment return segment

View File

@ -1,7 +1,7 @@
import numpy as np import numpy as np
from errors import * from errors import *
from segments import EmptySegment, ObjSegment, RawSectorsSegment, IndexedByteSegment from segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment
from utils import to_numpy from utils import to_numpy
class AtrHeader(object): class AtrHeader(object):
@ -92,18 +92,11 @@ class XfdHeader(AtrHeader):
class DiskImageBase(object): class DiskImageBase(object):
debug = False def __init__(self, rawdata, filename=""):
self.rawdata = rawdata
def __init__(self, bytes, style=None, filename=""): self.bytes = self.rawdata.get_data()
self.bytes = to_numpy(bytes) self.style = self.rawdata.get_style()
self.size = np.alen(self.bytes) self.size = np.alen(self.bytes)
if style is None:
if self.debug:
self.style = np.arange(self.size, dtype=np.uint8)
else:
self.style = np.zeros(self.size, dtype=np.uint8)
else:
self.style = style
self.set_filename(filename) self.set_filename(filename)
self.header = None self.header = None
self.total_sectors = 0 self.total_sectors = 0
@ -164,7 +157,7 @@ class DiskImageBase(object):
pos, size = self.header.get_pos(sector) pos, size = self.header.get_pos(sector)
return self.bytes[pos:pos + size], pos, size return self.bytes[pos:pos + size], pos, size
def get_sectors(self, start, end=None): def get_sector_slice(self, start, end=None):
""" Get contiguous sectors """ Get contiguous sectors
:param start: first sector number to read (note: numbering starts from 1) :param start: first sector number to read (note: numbering starts from 1)
@ -178,7 +171,17 @@ class DiskImageBase(object):
start += 1 start += 1
_, more = self.header.get_pos(start) _, more = self.header.get_pos(start)
size += more size += more
return self.bytes[pos:pos + size], self.style[pos:pos + size] return slice(pos, pos + size)
def get_sectors(self, start, end=None):
""" Get contiguous sectors
:param start: first sector number to read (note: numbering starts from 1)
:param end: last sector number to read
:returns: bytes
"""
s = self.get_sector_slice(start, end)
return self.bytes[s], self.style[s]
def get_contiguous_sectors(self, sector, num): def get_contiguous_sectors(self, sector, num):
start = 0 start = 0
@ -191,12 +194,11 @@ class DiskImageBase(object):
return start, count return start, count
def parse_segments(self): def parse_segments(self):
b = self.bytes r = self.rawdata
s = self.style
i = self.header.atr_header_offset i = self.header.atr_header_offset
if self.header.image_size > 0: if self.header.image_size > 0:
self.segments.append(ObjSegment(b[0:i], s[0:i], 0, 0, 0, i, name="%s Header" % self.header.file_format)) self.segments.append(ObjSegment(r[0:i], 0, 0, 0, i, name="%s Header" % self.header.file_format))
self.segments.append(RawSectorsSegment(b[i:], s[i:], 1, self.header.max_sectors, self.header.image_size, 128, 3, self.header.sector_size, name="Raw disk sectors")) 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"))
self.segments.extend(self.get_boot_segments()) self.segments.extend(self.get_boot_segments())
self.segments.extend(self.get_vtoc_segments()) self.segments.extend(self.get_vtoc_segments())
self.segments.extend(self.get_directory_segments()) self.segments.extend(self.get_directory_segments())
@ -217,10 +219,11 @@ class DiskImageBase(object):
if flag == 0: if flag == 0:
num = int(values[1]) num = int(values[1])
addr = int(values[2]) addr = int(values[2])
bytes, style = self.get_sectors(1, num) s = self.get_sector_slice(1, num)
header = ObjSegment(bytes[0:6], style[0:6], 0, 0, addr, addr + 6, name="Boot Header") r = self.rawdata[s]
sectors = ObjSegment(bytes, style, 0, 0, addr, addr + len(bytes), name="Boot Sectors") header = ObjSegment(r[0:6], 0, 0, addr, addr + 6, name="Boot Header")
code = ObjSegment(bytes[6:], style[6:], 0, 0, addr + 6, addr + len(bytes), name="Boot Code") sectors = ObjSegment(r, 0, 0, addr, addr + len(r), name="Boot Sectors")
code = ObjSegment(r[6:], 0, 0, addr + 6, addr + len(r), name="Boot Code")
segments = [sectors, header, code] segments = [sectors, header, code]
return segments return segments
@ -249,7 +252,7 @@ class DiskImageBase(object):
try: try:
segment = self.get_file_segment(dirent) segment = self.get_file_segment(dirent)
except InvalidFile, e: except InvalidFile, e:
segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename(), error=str(e)) segment = EmptySegment(self.rawdata, name=dirent.get_filename(), error=str(e))
segments.append(segment) segments.append(segment)
return segments return segments

View File

@ -1,7 +1,7 @@
import numpy as np import numpy as np
from errors import * from errors import *
from utils import to_numpy_list from utils import to_numpy, to_numpy_list
class SegmentSaver(object): class SegmentSaver(object):
@ -22,13 +22,47 @@ class SegmentSaver(object):
return "|".join(wildcards) return "|".join(wildcards)
class SegmentData(object):
def __init__(self, data, style=None, comments=None, debug=False):
self.data = to_numpy(data)
if style is None:
if debug:
self.style = np.arange(len(self), dtype=np.uint8)
else:
self.style = np.zeros(len(self), dtype=np.uint8)
else:
self.style = style
if comments is None:
comments = dict()
self.comments = comments
def __len__(self):
return np.alen(self.data)
def get_data(self):
return self.data
def get_style(self):
return self.style
def get_comments(self):
return self.comments
def __getitem__(self, index):
d = self.data[index]
s = self.style[index]
c = self.comments
return SegmentData(d, s, c)
class DefaultSegment(object): class DefaultSegment(object):
savers = [SegmentSaver] savers = [SegmentSaver]
def __init__(self, data, style, start_addr=0, name="All", error=None, verbose_name=None): def __init__(self, rawdata, start_addr=0, name="All", error=None, verbose_name=None):
self.start_addr = int(start_addr) # force python int to decouple from possibly being a numpy datatype self.start_addr = int(start_addr) # force python int to decouple from possibly being a numpy datatype
self.data = data self.rawdata = rawdata
self.style = style self.data = rawdata.get_data()
self.style = rawdata.get_style()
self.error = error self.error = error
self.name = name self.name = name
self.verbose_name = verbose_name self.verbose_name = verbose_name
@ -51,7 +85,7 @@ class DefaultSegment(object):
return s return s
def __len__(self): def __len__(self):
return np.alen(self.data) return len(self.rawdata)
def __getitem__(self, index): def __getitem__(self, index):
return self.data[index] return self.data[index]
@ -213,8 +247,8 @@ class DefaultSegment(object):
class EmptySegment(DefaultSegment): class EmptySegment(DefaultSegment):
def __init__(self, data, style, name="", error=None): def __init__(self, rawdata, name="", error=None):
DefaultSegment.__init__(self, data, style, 0, name, error) DefaultSegment.__init__(self, rawdata, 0, name, error)
def __str__(self): def __str__(self):
s = "%s (empty file)" % (self.name, ) s = "%s (empty file)" % (self.name, )
@ -234,8 +268,8 @@ class EmptySegment(DefaultSegment):
class ObjSegment(DefaultSegment): class ObjSegment(DefaultSegment):
def __init__(self, data, style, metadata_start, data_start, start_addr, end_addr, name="", **kwargs): def __init__(self, rawdata, metadata_start, data_start, start_addr, end_addr, name="", **kwargs):
DefaultSegment.__init__(self, data, style, start_addr, name, **kwargs) DefaultSegment.__init__(self, rawdata, start_addr, name, **kwargs)
self.metadata_start = metadata_start self.metadata_start = metadata_start
self.data_start = data_start self.data_start = data_start
@ -257,8 +291,8 @@ class ObjSegment(DefaultSegment):
class RawSectorsSegment(DefaultSegment): class RawSectorsSegment(DefaultSegment):
def __init__(self, data, style, first_sector, num_sectors, count, boot_sector_size, num_boot_sectors, sector_size, **kwargs): def __init__(self, rawdata, first_sector, num_sectors, count, boot_sector_size, num_boot_sectors, sector_size, **kwargs):
DefaultSegment.__init__(self, data, style, 0, **kwargs) DefaultSegment.__init__(self, rawdata, 0, **kwargs)
self.boot_sector_size = boot_sector_size self.boot_sector_size = boot_sector_size
self.num_boot_sectors = num_boot_sectors self.num_boot_sectors = num_boot_sectors
self.page_size = sector_size self.page_size = sector_size
@ -324,10 +358,11 @@ class IndexedStyleWrapper(object):
class IndexedByteSegment(DefaultSegment): class IndexedByteSegment(DefaultSegment):
def __init__(self, data, style, byte_order, **kwargs): def __init__(self, rawdata, byte_order, **kwargs):
# Convert to numpy list so fancy indexing works as argument to __getitem__ # Convert to numpy list so fancy indexing works as argument to __getitem__
self.order = to_numpy_list(byte_order) self.order = to_numpy_list(byte_order)
DefaultSegment.__init__(self, data, IndexedStyleWrapper(style, byte_order), 0, **kwargs) DefaultSegment.__init__(self, rawdata, **kwargs)
self.style = IndexedStyleWrapper(self.style, byte_order)
def __str__(self): def __str__(self):
s = "%s ($%x @ $%x)" % (self.name, len(self), self.order[0]) s = "%s ($%x @ $%x)" % (self.name, len(self), self.order[0])

View File

@ -179,19 +179,19 @@ class SpartaDosDiskImage(DiskImageBase):
def get_boot_segments(self): def get_boot_segments(self):
segments = [] segments = []
num = min(self.num_boot, 1) num = min(self.num_boot, 1)
bytes, style = self.get_sectors(1, num) s = self.get_sector_slice(1, num)
r = self.rawdata[s]
addr = self.boot_addr addr = self.boot_addr
header = ObjSegment(bytes[0:43], style[0:43], 0, 0, addr, addr + 43, name="Boot Header") header = ObjSegment(r[0:43], 0, 0, addr, addr + 43, name="Boot Header")
segments.append(header) segments.append(header)
if self.num_boot > 0: if self.num_boot > 0:
sectors = ObjSegment(bytes, style, 0, 0, addr, addr + len(bytes), name="Boot Sectors") sectors = ObjSegment(r, 0, 0, addr, addr + len(r), name="Boot Sectors")
code = ObjSegment(bytes[43:], style[43:], 0, 0, addr + 43, addr + len(bytes), name="Boot Code") code = ObjSegment(r[43:], 0, 0, addr + 43, addr + len(r), name="Boot Code")
segments.extend([header, code]) segments.extend([header, code])
return segments return segments
def get_vtoc_segments(self): def get_vtoc_segments(self):
b = self.bytes r = self.rawdata
s = self.style
segments = [] segments = []
addr = 0 addr = 0
start, count = self.get_contiguous_sectors(self.first_bitmap, self.num_bitmap) start, count = self.get_contiguous_sectors(self.first_bitmap, self.num_bitmap)
@ -201,7 +201,7 @@ class SpartaDosDiskImage(DiskImageBase):
else: else:
num_boot = 3 num_boot = 3
boot_size = 128 boot_size = 128
segment = RawSectorsSegment(b[start:start+count], s[start:start+count], self.first_bitmap, self.num_bitmap, count, 0, 0, self.sector_size, name="Bitmap") segment = RawSectorsSegment(r[start:start+count], self.first_bitmap, self.num_bitmap, count, 0, 0, self.sector_size, name="Bitmap")
segments.append(segment) segments.append(segment)
return segments return segments
@ -236,7 +236,7 @@ class SpartaDosDiskImage(DiskImageBase):
if len(byte_order) > 0: if len(byte_order) > 0:
name = "%s %d@%d %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.str_timestamp) name = "%s %d@%d %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.str_timestamp)
verbose_name = "%s (%d bytes, sector map@%d) %s %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.verbose_info, dirent.str_timestamp) verbose_name = "%s (%d bytes, sector map@%d) %s %s" % (dirent.get_filename(), dirent.length, dirent.starting_sector, dirent.verbose_info, dirent.str_timestamp)
segment = IndexedByteSegment(self.bytes, self.style, byte_order, name=name, verbose_name=verbose_name) segment = IndexedByteSegment(self.rawdata, byte_order, name=name, verbose_name=verbose_name)
else: else:
segment = EmptySegment(self.bytes, self.style, name=dirent.get_filename(), error=dirent.str_timestamp) segment = EmptySegment(self.rawdata, name=dirent.get_filename(), error=dirent.str_timestamp)
return segment return segment