xref: /openbmc/qemu/tests/qemu-iotests/qcow2.py (revision 4305d482)
1#!/usr/bin/env python
2
3from __future__ import print_function
4import sys
5import struct
6import string
7
8class QcowHeaderExtension:
9
10    def __init__(self, magic, length, data):
11        if length % 8 != 0:
12            padding = 8 - (length % 8)
13            data += b"\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
23class 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, 'mask', 'incompatible_features' ],
46        [ uint64_t, 'mask', 'compatible_features' ],
47        [ uint64_t, 'mask', '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, b""))
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            value = self.__dict__[f[2]]
134            if f[1] == 'mask':
135                bits = []
136                for bit in range(64):
137                    if value & (1 << bit):
138                        bits.append(bit)
139                value_str = str(bits)
140            else:
141                value_str = f[1] % value
142
143            print("%-25s" % f[2], value_str)
144        print("")
145
146    def dump_extensions(self):
147        for ex in self.extensions:
148
149            data = ex.data[:ex.length]
150            if all(c in string.printable.encode('ascii') for c in data):
151                data = "'%s'" % data.decode('ascii')
152            else:
153                data = "<binary>"
154
155            print("Header extension:")
156            print("%-25s %#x" % ("magic", ex.magic))
157            print("%-25s %d" % ("length", ex.length))
158            print("%-25s %s" % ("data", data))
159            print("")
160
161
162def cmd_dump_header(fd):
163    h = QcowHeader(fd)
164    h.dump()
165    h.dump_extensions()
166
167def cmd_dump_header_exts(fd):
168    h = QcowHeader(fd)
169    h.dump_extensions()
170
171def cmd_set_header(fd, name, value):
172    try:
173        value = int(value, 0)
174    except:
175        print("'%s' is not a valid number" % value)
176        sys.exit(1)
177
178    fields = (field[2] for field in QcowHeader.fields)
179    if not name in fields:
180        print("'%s' is not a known header field" % name)
181        sys.exit(1)
182
183    h = QcowHeader(fd)
184    h.__dict__[name] = value
185    h.update(fd)
186
187def cmd_add_header_ext(fd, magic, data):
188    try:
189        magic = int(magic, 0)
190    except:
191        print("'%s' is not a valid magic number" % magic)
192        sys.exit(1)
193
194    h = QcowHeader(fd)
195    h.extensions.append(QcowHeaderExtension.create(magic, data.encode('ascii')))
196    h.update(fd)
197
198def cmd_add_header_ext_stdio(fd, magic):
199    data = sys.stdin.read()
200    cmd_add_header_ext(fd, magic, data)
201
202def cmd_del_header_ext(fd, magic):
203    try:
204        magic = int(magic, 0)
205    except:
206        print("'%s' is not a valid magic number" % magic)
207        sys.exit(1)
208
209    h = QcowHeader(fd)
210    found = False
211
212    for ex in h.extensions:
213        if ex.magic == magic:
214            found = True
215            h.extensions.remove(ex)
216
217    if not found:
218        print("No such header extension")
219        return
220
221    h.update(fd)
222
223def cmd_set_feature_bit(fd, group, bit):
224    try:
225        bit = int(bit, 0)
226        if bit < 0 or bit >= 64:
227            raise ValueError
228    except:
229        print("'%s' is not a valid bit number in range [0, 64)" % bit)
230        sys.exit(1)
231
232    h = QcowHeader(fd)
233    if group == 'incompatible':
234        h.incompatible_features |= 1 << bit
235    elif group == 'compatible':
236        h.compatible_features |= 1 << bit
237    elif group == 'autoclear':
238        h.autoclear_features |= 1 << bit
239    else:
240        print("'%s' is not a valid group, try 'incompatible', 'compatible', or 'autoclear'" % group)
241        sys.exit(1)
242
243    h.update(fd)
244
245cmds = [
246    [ 'dump-header',          cmd_dump_header,          0, 'Dump image header and header extensions' ],
247    [ 'dump-header-exts',     cmd_dump_header_exts,     0, 'Dump image header extensions' ],
248    [ 'set-header',           cmd_set_header,           2, 'Set a field in the header'],
249    [ 'add-header-ext',       cmd_add_header_ext,       2, 'Add a header extension' ],
250    [ 'add-header-ext-stdio', cmd_add_header_ext_stdio, 1, 'Add a header extension, data from stdin' ],
251    [ 'del-header-ext',       cmd_del_header_ext,       1, 'Delete a header extension' ],
252    [ 'set-feature-bit',      cmd_set_feature_bit,      2, 'Set a feature bit'],
253]
254
255def main(filename, cmd, args):
256    fd = open(filename, "r+b")
257    try:
258        for name, handler, num_args, desc in cmds:
259            if name != cmd:
260                continue
261            elif len(args) != num_args:
262                usage()
263                return
264            else:
265                handler(fd, *args)
266                return
267        print("Unknown command '%s'" % cmd)
268    finally:
269        fd.close()
270
271def usage():
272    print("Usage: %s <file> <cmd> [<arg>, ...]" % sys.argv[0])
273    print("")
274    print("Supported commands:")
275    for name, handler, num_args, desc in cmds:
276        print("    %-20s - %s" % (name, desc))
277
278if __name__ == '__main__':
279    if len(sys.argv) < 3:
280        usage()
281        sys.exit(1)
282
283    main(sys.argv[1], sys.argv[2], sys.argv[3:])
284