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