tests/qemu-iotests/qcow2.py

   1 #!/usr/bin/env python
   2
   3 from __future__ import print_function
   4 import sys
   5 import struct
   6 import string
   7
   8 class QcowHeaderExtension:
   9
  10     def __init__(self, magic, length, data):
  11         if length % 8 != 0:
  12             padding = 8 - (length % 8)
  13             data += "\0" * padding
  14
  15         self.magic  = magic
  16         self.length = length
  17         self.data   = data
  18
  19     @classmethod
  20     def create(cls, magic, data):
  21         return QcowHeaderExtension(magic, len(data), data)
  22
  23 class QcowHeader:
  24
  25     uint32_t = 'I'
  26     uint64_t = 'Q'
  27
  28     fields = [
  29         # Version 2 header fields
  30         [ uint32_t, '%#x',  'magic' ],
  31         [ uint32_t, '%d',   'version' ],
  32         [ uint64_t, '%#x',  'backing_file_offset' ],
  33         [ uint32_t, '%#x',  'backing_file_size' ],
  34         [ uint32_t, '%d',   'cluster_bits' ],
  35         [ uint64_t, '%d',   'size' ],
  36         [ uint32_t, '%d',   'crypt_method' ],
  37         [ uint32_t, '%d',   'l1_size' ],
  38         [ uint64_t, '%#x',  'l1_table_offset' ],
  39         [ uint64_t, '%#x',  'refcount_table_offset' ],
  40         [ uint32_t, '%d',   'refcount_table_clusters' ],
  41         [ uint32_t, '%d',   'nb_snapshots' ],
  42         [ uint64_t, '%#x',  'snapshot_offset' ],
  43
  44         # Version 3 header fields
  45         [ uint64_t, '%#x',  'incompatible_features' ],
  46         [ uint64_t, '%#x',  'compatible_features' ],
  47         [ uint64_t, '%#x',  'autoclear_features' ],
  48         [ uint32_t, '%d',   'refcount_order' ],
  49         [ uint32_t, '%d',   'header_length' ],
  50     ];
  51
  52     fmt = '>' + ''.join(field[0] for field in fields)
  53
  54     def __init__(self, fd):
  55
  56         buf_size = struct.calcsize(QcowHeader.fmt)
  57
  58         fd.seek(0)
  59         buf = fd.read(buf_size)
  60
  61         header = struct.unpack(QcowHeader.fmt, buf)
  62         self.__dict__ = dict((field[2], header[i])
  63             for i, field in enumerate(QcowHeader.fields))
  64
  65         self.set_defaults()
  66         self.cluster_size = 1 << self.cluster_bits
  67
  68         fd.seek(self.header_length)
  69         self.load_extensions(fd)
  70
  71         if self.backing_file_offset:
  72             fd.seek(self.backing_file_offset)
  73             self.backing_file = fd.read(self.backing_file_size)
  74         else:
  75             self.backing_file = None
  76
  77     def set_defaults(self):
  78         if self.version == 2:
  79             self.incompatible_features = 0
  80             self.compatible_features = 0
  81             self.autoclear_features = 0
  82             self.refcount_order = 4
  83             self.header_length = 72
  84
  85     def load_extensions(self, fd):
  86         self.extensions = []
  87
  88         if self.backing_file_offset != 0:
  89             end = min(self.cluster_size, self.backing_file_offset)
  90         else:
  91             end = self.cluster_size
  92
  93         while fd.tell() < end:
  94             (magic, length) = struct.unpack('>II', fd.read(8))
  95             if magic == 0:
  96                 break
  97             else:
  98                 padded = (length + 7) & ~7
  99                 data = fd.read(padded)
 100                 self.extensions.append(QcowHeaderExtension(magic, length, data))
 101
 102     def update_extensions(self, fd):
 103
 104         fd.seek(self.header_length)
 105         extensions = self.extensions
 106         extensions.append(QcowHeaderExtension(0, 0, ""))
 107         for ex in extensions:
 108             buf = struct.pack('>II', ex.magic, ex.length)
 109             fd.write(buf)
 110             fd.write(ex.data)
 111
 112         if self.backing_file != None:
 113             self.backing_file_offset = fd.tell()
 114             fd.write(self.backing_file)
 115
 116         if fd.tell() > self.cluster_size:
 117             raise Exception("I think I just broke the image...")
 118
 119
 120     def update(self, fd):
 121         header_bytes = self.header_length
 122
 123         self.update_extensions(fd)
 124
 125         fd.seek(0)
 126         header = tuple(self.__dict__[f] for t, p, f in QcowHeader.fields)
 127         buf = struct.pack(QcowHeader.fmt, *header)
 128         buf = buf[0:header_bytes-1]
 129         fd.write(buf)
 130
 131     def dump(self):
 132         for f in QcowHeader.fields:
 133             print("%-25s" % f[2], f[1] % self.__dict__[f[2]])
 134         print("")
 135
 136     def dump_extensions(self):
 137         for ex in self.extensions:
 138
 139             data = ex.data[:ex.length]
 140             if all(c in string.printable for c in data):
 141                 data = "'%s'" % data
 142             else:
 143                 data = "<binary>"
 144
 145             print("Header extension:")
 146             print("%-25s %#x" % ("magic", ex.magic))
 147             print("%-25s %d" % ("length", ex.length))
 148             print("%-25s %s" % ("data", data))
 149             print("")
 150
 151
 152 def cmd_dump_header(fd):
 153     h = QcowHeader(fd)
 154     h.dump()
 155     h.dump_extensions()
 156
 157 def cmd_set_header(fd, name, value):
 158     try:
 159         value = int(value, 0)
 160     except:
 161         print("'%s' is not a valid number" % value)
 162         sys.exit(1)
 163
 164     fields = (field[2] for field in QcowHeader.fields)
 165     if not name in fields:
 166         print("'%s' is not a known header field" % name)
 167         sys.exit(1)
 168
 169     h = QcowHeader(fd)
 170     h.__dict__[name] = value
 171     h.update(fd)
 172
 173 def cmd_add_header_ext(fd, magic, data):
 174     try:
 175         magic = int(magic, 0)
 176     except:
 177         print("'%s' is not a valid magic number" % magic)
 178         sys.exit(1)
 179
 180     h = QcowHeader(fd)
 181     h.extensions.append(QcowHeaderExtension.create(magic, data))
 182     h.update(fd)
 183
 184 def cmd_add_header_ext_stdio(fd, magic):
 185     data = sys.stdin.read()
 186     cmd_add_header_ext(fd, magic, data)
 187
 188 def cmd_del_header_ext(fd, magic):
 189     try:
 190         magic = int(magic, 0)
 191     except:
 192         print("'%s' is not a valid magic number" % magic)
 193         sys.exit(1)
 194
 195     h = QcowHeader(fd)
 196     found = False
 197
 198     for ex in h.extensions:
 199         if ex.magic == magic:
 200             found = True
 201             h.extensions.remove(ex)
 202
 203     if not found:
 204         print("No such header extension")
 205         return
 206
 207     h.update(fd)
 208
 209 def cmd_set_feature_bit(fd, group, bit):
 210     try:
 211         bit = int(bit, 0)
 212         if bit < 0 or bit >= 64:
 213             raise ValueError
 214     except:
 215         print("'%s' is not a valid bit number in range [0, 64)" % bit)
 216         sys.exit(1)
 217
 218     h = QcowHeader(fd)
 219     if group == 'incompatible':
 220         h.incompatible_features |= 1 << bit
 221     elif group == 'compatible':
 222         h.compatible_features |= 1 << bit
 223     elif group == 'autoclear':
 224         h.autoclear_features |= 1 << bit
 225     else:
 226         print("'%s' is not a valid group, try 'incompatible', 'compatible', or 'autoclear'" % group)
 227         sys.exit(1)
 228
 229     h.update(fd)
 230
 231 cmds = [
 232     [ 'dump-header',          cmd_dump_header,          0, 'Dump image header and header extensions' ],
 233     [ 'set-header',           cmd_set_header,           2, 'Set a field in the header'],
 234     [ 'add-header-ext',       cmd_add_header_ext,       2, 'Add a header extension' ],
 235     [ 'add-header-ext-stdio', cmd_add_header_ext_stdio, 1, 'Add a header extension, data from stdin' ],
 236     [ 'del-header-ext',       cmd_del_header_ext,       1, 'Delete a header extension' ],
 237     [ 'set-feature-bit',      cmd_set_feature_bit,      2, 'Set a feature bit'],
 238 ]
 239
 240 def main(filename, cmd, args):
 241     fd = open(filename, "r+b")
 242     try:
 243         for name, handler, num_args, desc in cmds:
 244             if name != cmd:
 245                 continue
 246             elif len(args) != num_args:
 247                 usage()
 248                 return
 249             else:
 250                 handler(fd, *args)
 251                 return
 252         print("Unknown command '%s'" % cmd)
 253     finally:
 254         fd.close()
 255
 256 def usage():
 257     print("Usage: %s <file> <cmd> [<arg>, ...]" % sys.argv[0])
 258     print("")
 259     print("Supported commands:")
 260     for name, handler, num_args, desc in cmds:
 261         print("    %-20s - %s" % (name, desc))
 262
 263 if __name__ == '__main__':
 264     if len(sys.argv) < 3:
 265         usage()
 266         sys.exit(1)
 267
 268     main(sys.argv[1], sys.argv[2], sys.argv[3:])