xref: /openbmc/qemu/tests/qemu-iotests/qcow2.py (revision 5964ed56)
1#!/usr/bin/env python3
2
3import sys
4import struct
5import string
6
7class 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
22class 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
161def cmd_dump_header(fd):
162    h = QcowHeader(fd)
163    h.dump()
164    h.dump_extensions()
165
166def cmd_dump_header_exts(fd):
167    h = QcowHeader(fd)
168    h.dump_extensions()
169
170def 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
186def 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
197def cmd_add_header_ext_stdio(fd, magic):
198    data = sys.stdin.read()
199    cmd_add_header_ext(fd, magic, data)
200
201def 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
222def 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
244cmds = [
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
254def 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
270def 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
277if __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:])
283