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