2016-02-13 04:36:33 +00:00
import numpy as np
from errors import *
2016-03-25 23:07:16 +00:00
from segments import SegmentData , EmptySegment , ObjSegment , RawSectorsSegment
2017-02-23 21:53:32 +00:00
from utils import *
2016-02-13 04:36:33 +00:00
2017-02-22 03:25:47 +00:00
import logging
log = logging . getLogger ( __name__ )
2017-02-23 21:02:56 +00:00
class BaseHeader ( object ) :
file_format = " generic " # text descriptor of file format
2017-02-23 21:53:32 +00:00
sector_class = WriteableSector
2017-02-23 21:02:56 +00:00
def __init__ ( self , sector_size = 256 , initial_sectors = 0 , vtoc_sector = 0 , starting_sector_label = 0 ) :
self . image_size = 0
self . sector_size = sector_size
2017-02-24 03:13:48 +00:00
self . payload_bytes = sector_size
2017-02-23 21:02:56 +00:00
self . initial_sector_size = 0
self . num_initial_sectors = 0
self . crc = 0
self . unused = 0
self . flags = 0
self . header_offset = 0
self . starting_sector_label = starting_sector_label
self . max_sectors = 0 # number of sectors, -1 is unlimited
self . tracks_per_disk = 0
self . sectors_per_track = 0
self . first_vtoc = vtoc_sector
self . num_vtoc = 1
self . extra_vtoc = [ ]
self . first_directory = 0
self . num_directory = 0
2017-03-23 17:06:37 +00:00
2017-02-23 21:02:56 +00:00
def __len__ ( self ) :
return self . header_offset
2017-03-23 17:06:37 +00:00
2017-02-23 21:02:56 +00:00
def to_array ( self ) :
header_bytes = np . zeros ( [ self . header_offset ] , dtype = np . uint8 )
self . encode ( header_bytes )
return header_bytes
def encode ( self , header_bytes ) :
""" Subclasses should override this to put the byte values into the
header .
"""
return
def sector_is_valid ( self , sector ) :
return ( self . max_sectors < 0 ) | ( sector > = self . starting_sector_label and sector < ( self . max_sectors + self . starting_sector_label ) )
2017-02-26 21:07:36 +00:00
def iter_sectors ( self ) :
i = self . starting_sector_label
while self . sector_is_valid ( i ) :
pos , size = self . get_pos ( i )
yield i , pos , size
i + = 1
2017-03-23 17:06:37 +00:00
2017-02-23 21:02:56 +00:00
def get_pos ( self , sector ) :
""" Get index (into the raw data of the disk image) of start of sector
This base class method assumes the sectors are one after another , in
order starting from the beginning of the raw data .
"""
if not self . sector_is_valid ( sector ) :
raise ByteNotInFile166 ( " Sector %d out of range " % sector )
pos = sector * self . sector_size + self . header_offset
size = self . sector_size
return pos , size
def sector_from_track ( self , track , sector ) :
return ( track * self . sectors_per_track ) + sector
def track_from_sector ( self , sector ) :
track , sector = divmod ( sector , self . sectors_per_track )
return track , sector
def check_size ( self , size ) :
raise InvalidDiskImage ( " BaseHeader subclasses need custom checks for size " )
def strict_check ( self , image ) :
pass
2017-02-23 21:53:32 +00:00
def create_sector ( self , data ) :
return self . sector_class ( self . sector_size , data )
2017-02-23 21:02:56 +00:00
2016-02-13 04:36:33 +00:00
class DiskImageBase ( object ) :
2016-03-25 23:07:16 +00:00
def __init__ ( self , rawdata , filename = " " ) :
self . rawdata = rawdata
self . bytes = self . rawdata . get_data ( )
self . style = self . rawdata . get_style ( )
2016-02-13 04:36:33 +00:00
self . size = np . alen ( self . bytes )
self . set_filename ( filename )
self . header = None
self . total_sectors = 0
self . unused_sectors = 0
2017-02-22 03:25:47 +00:00
self . files = [ ] # all dirents that show up in a normal dir listing
2016-02-13 04:36:33 +00:00
self . segments = [ ]
self . all_sane = True
2017-02-24 19:57:51 +00:00
self . default_filetype = " "
2016-02-13 04:36:33 +00:00
self . setup ( )
2016-06-02 22:09:34 +00:00
def __len__ ( self ) :
return len ( self . rawdata )
2017-02-22 03:25:47 +00:00
@property
def writeable_sector_class ( self ) :
return WriteableSector
2017-02-22 20:12:52 +00:00
@property
def raw_sector_class ( self ) :
return RawSectorsSegment
2017-02-22 03:25:47 +00:00
@property
def vtoc_class ( self ) :
return VTOC
@property
def directory_class ( self ) :
return Directory
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def set_filename ( self , filename ) :
if " . " in filename :
self . filename , self . ext = filename . rsplit ( " . " , 1 )
else :
self . filename , self . ext = filename , " "
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def dir ( self ) :
lines = [ ]
lines . append ( str ( self ) )
for dirent in self . files :
if dirent . in_use :
lines . append ( str ( dirent ) )
return " \n " . join ( lines )
def setup ( self ) :
self . size = np . alen ( self . bytes )
2016-06-01 21:54:52 +00:00
self . read_header ( )
2016-05-06 06:11:06 +00:00
self . header . check_size ( self . size - len ( self . header ) )
2016-02-13 04:36:33 +00:00
self . check_size ( )
2017-02-22 03:25:47 +00:00
self . get_metadata ( )
def get_metadata ( self ) :
2016-02-13 04:36:33 +00:00
self . get_boot_sector_info ( )
self . get_vtoc ( )
self . get_directory ( )
self . check_sane ( )
2016-06-02 23:59:59 +00:00
def strict_check ( self ) :
""" Perform the strictest of checks to verify the data is valid """
self . header . strict_check ( self )
def relaxed_check ( self ) :
""" Conform as much as possible to get the data to work with this
format .
"""
pass
2017-03-23 17:06:37 +00:00
2016-05-06 06:11:06 +00:00
@classmethod
def new_header ( cls , diskimage , format = " ATR " ) :
2017-02-23 22:23:29 +00:00
raise NotImplementedError
2017-03-23 17:06:37 +00:00
2016-05-06 06:11:06 +00:00
def as_new_format ( self , format = " ATR " ) :
""" Create a new disk image in the specified format
"""
2017-02-23 22:23:29 +00:00
raise NotImplementedError
2017-03-23 17:06:37 +00:00
2016-05-05 20:59:56 +00:00
def save ( self , filename = " " ) :
if not filename :
filename = self . filename
if self . ext :
filename + = " . " + self . ext
if not filename :
raise RuntimeError ( " No filename specified for save! " )
bytes = self . bytes [ : ]
with open ( filename , " wb " ) as fh :
bytes . tofile ( fh )
2017-03-23 17:06:37 +00:00
2016-03-28 22:48:46 +00:00
def assert_valid_sector ( self , sector ) :
if not self . header . sector_is_valid ( sector ) :
raise ByteNotInFile166 ( " Sector %d out of range " % sector )
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def check_sane ( self ) :
if not self . all_sane :
raise InvalidDiskImage ( " Invalid directory entries; may be boot disk " )
2017-03-23 17:06:37 +00:00
2016-06-01 21:54:52 +00:00
def read_header ( self ) :
2017-02-23 23:34:56 +00:00
return BaseHeader ( )
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def check_size ( self ) :
2016-03-02 22:05:32 +00:00
pass
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_boot_sector_info ( self ) :
pass
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_vtoc ( self ) :
2017-02-22 03:25:47 +00:00
""" Get information from VTOC and populate the VTOC object """
2016-02-13 04:36:33 +00:00
pass
2017-03-23 17:06:37 +00:00
2017-02-22 03:25:47 +00:00
def get_directory ( self , directory = None ) :
2016-02-13 04:36:33 +00:00
pass
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_raw_bytes ( self , sector ) :
pos , size = self . header . get_pos ( sector )
return self . bytes [ pos : pos + size ] , pos , size
2017-03-23 17:06:37 +00:00
2016-03-25 23:07:16 +00:00
def get_sector_slice ( self , start , end = None ) :
2016-02-13 04:36:33 +00:00
""" Get contiguous sectors
: param start : first sector number to read ( note : numbering starts from 1 )
: param end : last sector number to read
: returns : bytes
"""
pos , size = self . header . get_pos ( start )
if end is None :
end = start
while start < end :
start + = 1
_ , more = self . header . get_pos ( start )
size + = more
2016-03-25 23:07:16 +00:00
return slice ( pos , pos + size )
2017-03-23 17:06:37 +00:00
2016-03-25 23:07:16 +00:00
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 ]
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_contiguous_sectors ( self , sector , num ) :
start = 0
count = 0
for index in range ( sector , sector + num ) :
pos , size = self . header . get_pos ( index )
if start == 0 :
start = pos
count + = size
return start , count
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def parse_segments ( self ) :
2016-03-25 23:07:16 +00:00
r = self . rawdata
2016-06-01 21:54:52 +00:00
i = self . header . header_offset
2016-02-13 04:36:33 +00:00
if self . header . image_size > 0 :
2016-03-25 23:07:16 +00:00
self . segments . append ( ObjSegment ( r [ 0 : i ] , 0 , 0 , 0 , i , name = " %s Header " % self . header . file_format ) )
2017-02-22 20:12:52 +00:00
self . segments . append ( self . raw_sector_class ( r [ i : ] , self . header . starting_sector_label , self . header . max_sectors , self . header . image_size , self . header . initial_sector_size , self . header . num_initial_sectors , self . header . sector_size , name = " Raw disk sectors " ) )
2016-02-13 04:36:33 +00:00
self . segments . extend ( self . get_boot_segments ( ) )
self . segments . extend ( self . get_vtoc_segments ( ) )
self . segments . extend ( self . get_directory_segments ( ) )
self . segments . extend ( self . get_file_segments ( ) )
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_boot_segments ( self ) :
2017-03-02 20:28:05 +00:00
return [ ]
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_vtoc_segments ( self ) :
return [ ]
2017-02-22 03:25:47 +00:00
2016-02-13 04:36:33 +00:00
def get_directory_segments ( self ) :
return [ ]
2017-03-23 17:06:37 +00:00
2017-02-22 07:07:24 +00:00
def find_dirent ( self , filename ) :
2017-02-26 22:06:29 +00:00
# check if we've been passed a dirent instead of a filename
if hasattr ( filename , " filename " ) :
return filename
2016-02-13 04:36:33 +00:00
for dirent in self . files :
2017-02-26 23:26:40 +00:00
if filename == dirent . filename :
2017-02-22 07:07:24 +00:00
return dirent
raise FileNotFound ( " %s not found on disk " % filename )
2017-03-23 17:06:37 +00:00
2017-02-22 07:07:24 +00:00
def find_file ( self , filename ) :
dirent = self . find_dirent ( filename )
return self . get_file ( dirent )
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_file ( self , dirent ) :
segment = self . get_file_segment ( dirent )
return segment . tostring ( )
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_file_segment ( self , dirent ) :
pass
2017-03-23 17:06:37 +00:00
2016-02-13 04:36:33 +00:00
def get_file_segments ( self ) :
segments = [ ]
for dirent in self . files :
try :
segment = self . get_file_segment ( dirent )
except InvalidFile , e :
2017-02-24 16:42:04 +00:00
segment = EmptySegment ( self . rawdata , name = dirent . filename , error = str ( e ) )
2016-02-13 04:36:33 +00:00
segments . append ( segment )
return segments
2017-02-26 20:45:54 +00:00
def create_executable_file_image ( self , segments , run_addr = None ) :
2017-02-26 20:14:23 +00:00
raise NotImplementedError
2017-02-22 03:25:47 +00:00
# file writing methods
2017-02-22 13:42:40 +00:00
def begin_transaction ( self ) :
state = self . bytes [ : ] , self . style [ : ]
return state
def rollback_transaction ( self , state ) :
self . bytes [ : ] , self . style [ : ] = state
return
2017-02-26 22:14:25 +00:00
def get_vtoc_object ( self ) :
vtoc_segments = self . get_vtoc_segments ( )
vtoc = self . vtoc_class ( self . header , vtoc_segments )
return vtoc
2017-02-22 03:25:47 +00:00
def write_file ( self , filename , filetype , data ) :
""" Write data to a file on disk
This throws various exceptions on failures , for instance if there is
not enough space on disk or a free entry is not available in the
catalog .
"""
2017-02-22 13:42:40 +00:00
state = self . begin_transaction ( )
try :
2017-02-23 21:53:32 +00:00
directory = self . directory_class ( self . header )
2017-02-22 13:42:40 +00:00
self . get_directory ( directory )
dirent = directory . add_dirent ( filename , filetype )
data = to_numpy ( data )
2017-02-24 03:13:48 +00:00
sector_list = self . build_sectors ( data )
2017-02-26 22:14:25 +00:00
vtoc = self . get_vtoc_object ( )
2017-02-22 15:19:52 +00:00
directory . save_dirent ( self , dirent , vtoc , sector_list )
2017-02-22 13:42:40 +00:00
self . write_sector_list ( sector_list )
self . write_sector_list ( vtoc )
self . write_sector_list ( directory )
except AtrError :
self . rollback_transaction ( state )
raise
finally :
self . get_metadata ( )
2017-02-22 03:25:47 +00:00
2017-02-24 03:13:48 +00:00
def build_sectors ( self , data ) :
data = to_numpy ( data )
2017-02-27 03:55:12 +00:00
sectors = BaseSectorList ( self . header )
2017-02-24 03:13:48 +00:00
index = 0
while index < len ( data ) :
count = min ( self . header . payload_bytes , len ( data ) - index )
sector = self . header . create_sector ( data [ index : index + count ] )
sectors . append ( sector )
index + = count
return sectors
2017-02-22 03:25:47 +00:00
def write_sector_list ( self , sector_list ) :
for sector in sector_list :
pos , size = self . header . get_pos ( sector . sector_num )
2017-05-07 19:06:39 +00:00
if _dbg : log . debug ( " writing: %s at %d " % ( sector , pos ) )
2017-02-22 03:25:47 +00:00
self . bytes [ pos : pos + size ] = sector . data
2017-02-22 07:07:24 +00:00
def delete_file ( self , filename ) :
2017-02-22 13:42:40 +00:00
state = self . begin_transaction ( )
try :
2017-02-23 21:53:32 +00:00
directory = self . directory_class ( self . header )
2017-02-22 13:42:40 +00:00
self . get_directory ( directory )
dirent = directory . find_dirent ( filename )
2017-02-24 06:16:18 +00:00
sector_list = dirent . get_sectors_in_vtoc ( self )
2017-02-26 22:14:25 +00:00
vtoc = self . get_vtoc_object ( )
2017-02-22 15:19:52 +00:00
directory . remove_dirent ( self , dirent , vtoc , sector_list )
2017-02-22 13:42:40 +00:00
self . write_sector_list ( vtoc )
self . write_sector_list ( directory )
except AtrError :
self . rollback_transaction ( state )
raise
finally :
self . get_metadata ( )
2017-02-26 21:07:36 +00:00
def shred ( self , fill_value = 0 ) :
state = self . begin_transaction ( )
try :
2017-02-26 22:14:25 +00:00
vtoc = self . get_vtoc_object ( )
2017-02-26 21:34:07 +00:00
for sector_num , pos , size in vtoc . iter_free_sectors ( ) :
2017-05-07 19:06:39 +00:00
if _dbg : log . debug ( " shredding: sector %s at %d , fill value= %d " % ( sector_num , pos , fill_value ) )
2017-02-26 21:07:36 +00:00
self . bytes [ pos : pos + size ] = fill_value
except AtrError :
self . rollback_transaction ( state )
raise
finally :
self . get_metadata ( )