From 4c4da4bf3815ad3e5bbd3b11e2e1f589cbe52a83 Mon Sep 17 00:00:00 2001 From: Rob McMullen Date: Thu, 5 May 2016 23:11:06 -0700 Subject: [PATCH] Added add_atr_header and removed hack to insert a header into an already existing disk image * until I think of a better solution, forcing you to create an entirely new disk image --- atrcopy/__init__.py | 2 +- atrcopy/diskimages.py | 48 ++++++++++++++++++++++++++++++++++--------- test/test_kboot.py | 24 ++++++++++------------ 3 files changed, 50 insertions(+), 24 deletions(-) diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py index a83752b..c26e7f5 100644 --- a/atrcopy/__init__.py +++ b/atrcopy/__init__.py @@ -9,7 +9,7 @@ except ImportError: from errors import * from ataridos import AtariDosDiskImage, AtariDosFile, get_xex -from diskimages import AtrHeader, BootDiskImage +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 from spartados import SpartaDosDiskImage diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py index c4a2989..8b37ab2 100644 --- a/atrcopy/diskimages.py +++ b/atrcopy/diskimages.py @@ -17,7 +17,7 @@ class AtrHeader(object): ]) file_format = "ATR" - def __init__(self, bytes=None, sector_size=128, initial_sectors=3): + def __init__(self, bytes=None, sector_size=128, initial_sectors=3, create=False): self.image_size = 0 self.sector_size = sector_size self.crc = 0 @@ -27,6 +27,9 @@ class AtrHeader(object): self.initial_sector_size = sector_size self.num_initial_sectors = initial_sectors self.max_sectors = 0 + if create: + self.atr_header_offset = 16 + self.check_size(0) if bytes is None: return @@ -46,6 +49,9 @@ class AtrHeader(object): def __str__(self): 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 + def to_array(self): raw = np.zeros([16], dtype=np.uint8) values = raw.view(dtype=self.format)[0] @@ -104,6 +110,9 @@ class XfdHeader(AtrHeader): def __str__(self): return "%s Disk Image (size=%d (%dx%db)" % (self.file_format, self.image_size, self.max_sectors, self.sector_size) + def __len__(self): + return 0 + def to_array(self): raw = np.zeros([0], dtype=np.uint8) return raw @@ -141,17 +150,31 @@ class DiskImageBase(object): def setup(self): self.size = np.alen(self.bytes) self.read_atr_header() - self.header.check_size(self.size) + self.header.check_size(self.size - len(self.header)) self.check_size() self.get_boot_sector_info() self.get_vtoc() self.get_directory() self.check_sane() - def rebuild_header(self): - header = AtrHeader() - header.check_size(self.size) - self.header = header + @classmethod + def new_header(cls, diskimage, format="ATR"): + if format.lower() == "atr": + header = AtrHeader(create=True) + header.check_size(diskimage.size) + else: + raise RuntimeError("Unknown header type %s" % format) + return header + + def as_new_format(self, format="ATR"): + """ Create a new disk image in the specified format + """ + first_data = len(self.header) + raw = self.rawdata[first_data:] + data = add_atr_header(raw) + newraw = SegmentData(data) + image = self.__class__(newraw) + return image def save(self, filename=""): if not filename: @@ -162,10 +185,6 @@ class DiskImageBase(object): raise RuntimeError("No filename specified for save!") bytes = self.bytes[:] with open(filename, "wb") as fh: - h = self.header.to_array() - print h - if len(h) > 0: - h.tofile(fh) bytes.tofile(fh) def assert_valid_sector(self, sector): @@ -323,3 +342,12 @@ class BootDiskImage(DiskImageBase): max_sectors = max_size / self.header.sector_size if nsec > max_sectors: raise InvalidDiskImage("Number of boot sectors out of range") + +def add_atr_header(bytes): + header = AtrHeader(create=True) + header.check_size(len(bytes)) + hlen = len(header) + data = np.empty([hlen + len(bytes)], dtype=np.uint8) + data[0:hlen] = header.to_array() + data[hlen:] = bytes + return data diff --git a/test/test_kboot.py b/test/test_kboot.py index fbd2af8..4e5bba4 100644 --- a/test/test_kboot.py +++ b/test/test_kboot.py @@ -2,7 +2,7 @@ import os import numpy as np -from atrcopy import SegmentData, KBootImage, add_xexboot_header +from atrcopy import SegmentData, KBootImage, add_xexboot_header, add_atr_header class TestKbootHeader(object): @@ -12,20 +12,18 @@ class TestKbootHeader(object): def check_size(self, data): xex_size = len(data) bytes = add_xexboot_header(data) +# rawdata = SegmentData(bytes) +# size = len(rawdata) +# atr = KBootImage(rawdata) +# newatr = atr.as_new_format("ATR") + image_size = len(bytes) + bytes = add_atr_header(bytes) rawdata = SegmentData(bytes) - size = len(rawdata) - atr = KBootImage(rawdata) - atr.rebuild_header() - header_data = atr.header.to_array() - print header_data - newdata = np.append(header_data, atr.bytes) - newrawdata = SegmentData(newdata) - atr = KBootImage(newrawdata) - atr.rebuild_header() - image = atr.bytes + newatr = KBootImage(rawdata) + image = newatr.bytes print image[0:16] - paragraphs = size / 16 - print atr.header, paragraphs + paragraphs = image_size / 16 + print newatr.header, paragraphs assert int(image[2:4].view(dtype='