tests/qemu-iotests/qcow2.py

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