1import os 2import stat 3import amitools.util.BlkDevTools as BlkDevTools 4import zlib 5import io 6 7class ImageFile: 8 def __init__(self, file_name, read_only=False, block_bytes=512, fobj=None): 9 self.file_name = file_name 10 self.read_only = read_only 11 self.block_bytes = block_bytes 12 self.fobj = fobj 13 self.fh = None 14 self.size = 0 15 self.num_blocks = 0 16 17 def open(self): 18 # file obj? 19 if self.fobj is not None: 20 self.fh = self.fobj 21 # get size via seek 22 self.fobj.seek(0,2) # end of file 23 self.size = self.fobj.tell() 24 self.fobj.seek(0,0) # return to begin 25 self.num_blocks = self.size / self.block_bytes 26 # file name given 27 else: 28 # is readable? 29 if not os.access(self.file_name, os.R_OK): 30 raise IOError("Can't read from image file") 31 # is writeable? 32 if not os.access(self.file_name, os.W_OK): 33 self.read_only = True 34 # is it a block/char device? 35 st = os.stat(self.file_name) 36 mode = st.st_mode 37 if stat.S_ISBLK(mode) or stat.S_ISCHR(mode): 38 self.size = BlkDevTools.getblkdevsize(self.file_name) 39 else: 40 # get size and make sure its not empty 41 self.size = os.path.getsize(self.file_name) 42 if self.size == 0: 43 raise IOError("Empty image file detected!") 44 self.num_blocks = self.size / self.block_bytes 45 # open raw file 46 if self.read_only: 47 flags = "rb" 48 else: 49 flags = "r+b" 50 self.fh = io.open(self.file_name, flags) 51 52 def read_blk(self, blk_num): 53 if blk_num >= self.num_blocks: 54 raise IOError("Invalid image file block num: got %d but max is %d" % (blk_num, self.num_blocks)) 55 off = blk_num * self.block_bytes 56 if off != self.fh.tell(): 57 self.fh.seek(off, os.SEEK_SET) 58 num = self.block_bytes 59 data = self.fh.read(self.block_bytes) 60 return data 61 62 def write_blk(self, blk_num, data): 63 if self.read_only: 64 raise IOError("Can't write block: image file is read-only") 65 if blk_num >= self.num_blocks: 66 raise IOError("Invalid image file block num: got %d but max is %d" % (blk_num, self.num_blocks)) 67 if len(data) != self.block_bytes: 68 raise IOError("Invalid block size written: got %d but size is %d" % (len(data), self.block_bytes)) 69 off = blk_num * self.block_bytes 70 if off != self.fh.tell(): 71 self.fh.seek(off, os.SEEK_SET) 72 self.fh.write(data) 73 74 def flush(self): 75 self.fh.flush() 76 77 def close(self): 78 if self.fh != None: 79 self.fh.close() 80 self.fh = None 81 82 def create(self, num_blocks): 83 if self.read_only: 84 raise IOError("Can't create image file in read only mode") 85 block = '\0' * self.block_bytes 86 if self.fobj is not None: 87 for i in xrange(num_blocks): 88 self.fobj.write(block) 89 self.fobj.seek(0,0) 90 else: 91 fh = file(self.file_name, "wb") 92 for i in xrange(num_blocks): 93 fh.write(block) 94 fh.close() 95 96 97# --- mini test --- 98if __name__ == '__main__': 99 import sys 100 for a in sys.argv[1:]: 101 # read image 102 im = ImageFile(a) 103 im.open() 104 d = im.read_blk(0) 105 im.write_blk(0,d) 106 im.close() 107 # read fobj 108 fobj = file(a,"r+b") 109 im = ImageFile(a,fobj=fobj) 110 im.open() 111 d = im.read_blk(0) 112 im.write_blk(0,d) 113 im.close() 114