xref: /openbmc/qemu/tests/qemu-iotests/qed.py (revision 10df8ff1) 
1#!/usr/bin/env python
2#
3# Tool to manipulate QED image files
4#
5# Copyright (C) 2010 IBM, Corp.
6#
7# Authors:
8#  Stefan Hajnoczi <stefanha@linux.vnet.ibm.com>
9#
10# This work is licensed under the terms of the GNU GPL, version 2 or later.
11# See the COPYING file in the top-level directory.
12
13from __future__ import print_function
14import sys
15import struct
16import random
17import optparse
18
19# This can be used as a module
20__all__ = ['QED_F_NEED_CHECK', 'QED']
21
22QED_F_NEED_CHECK = 0x02
23
24header_fmt = '<IIIIQQQQQII'
25header_size = struct.calcsize(header_fmt)
26field_names = ['magic', 'cluster_size', 'table_size',
27               'header_size', 'features', 'compat_features',
28               'autoclear_features', 'l1_table_offset', 'image_size',
29               'backing_filename_offset', 'backing_filename_size']
30table_elem_fmt = '<Q'
31table_elem_size = struct.calcsize(table_elem_fmt)
32
33def err(msg):
34    sys.stderr.write(msg + '\n')
35    sys.exit(1)
36
37def unpack_header(s):
38    fields = struct.unpack(header_fmt, s)
39    return dict((field_names[idx], val) for idx, val in enumerate(fields))
40
41def pack_header(header):
42    fields = tuple(header[x] for x in field_names)
43    return struct.pack(header_fmt, *fields)
44
45def unpack_table_elem(s):
46    return struct.unpack(table_elem_fmt, s)[0]
47
48def pack_table_elem(elem):
49    return struct.pack(table_elem_fmt, elem)
50
51class QED(object):
52    def __init__(self, f):
53        self.f = f
54
55        self.f.seek(0, 2)
56        self.filesize = f.tell()
57
58        self.load_header()
59        self.load_l1_table()
60
61    def raw_pread(self, offset, size):
62        self.f.seek(offset)
63        return self.f.read(size)
64
65    def raw_pwrite(self, offset, data):
66        self.f.seek(offset)
67        return self.f.write(data)
68
69    def load_header(self):
70        self.header = unpack_header(self.raw_pread(0, header_size))
71
72    def store_header(self):
73        self.raw_pwrite(0, pack_header(self.header))
74
75    def read_table(self, offset):
76        size = self.header['table_size'] * self.header['cluster_size']
77        s = self.raw_pread(offset, size)
78        table = [unpack_table_elem(s[i:i + table_elem_size]) for i in xrange(0, size, table_elem_size)]
79        return table
80
81    def load_l1_table(self):
82        self.l1_table = self.read_table(self.header['l1_table_offset'])
83        self.table_nelems = self.header['table_size'] * self.header['cluster_size'] // table_elem_size
84
85    def write_table(self, offset, table):
86        s = ''.join(pack_table_elem(x) for x in table)
87        self.raw_pwrite(offset, s)
88
89def random_table_item(table):
90    vals = [(index, offset) for index, offset in enumerate(table) if offset != 0]
91    if not vals:
92        err('cannot pick random item because table is empty')
93    return random.choice(vals)
94
95def corrupt_table_duplicate(table):
96    '''Corrupt a table by introducing a duplicate offset'''
97    victim_idx, victim_val = random_table_item(table)
98    unique_vals = set(table)
99    if len(unique_vals) == 1:
100        err('no duplication corruption possible in table')
101    dup_val = random.choice(list(unique_vals.difference([victim_val])))
102    table[victim_idx] = dup_val
103
104def corrupt_table_invalidate(qed, table):
105    '''Corrupt a table by introducing an invalid offset'''
106    index, _ = random_table_item(table)
107    table[index] = qed.filesize + random.randint(0, 100 * 1024 * 1024 * 1024 * 1024)
108
109def cmd_show(qed, *args):
110    '''show [header|l1|l2 <offset>]- Show header or l1/l2 tables'''
111    if not args or args[0] == 'header':
112        print(qed.header)
113    elif args[0] == 'l1':
114        print(qed.l1_table)
115    elif len(args) == 2 and args[0] == 'l2':
116        offset = int(args[1])
117        print(qed.read_table(offset))
118    else:
119        err('unrecognized sub-command')
120
121def cmd_duplicate(qed, table_level):
122    '''duplicate l1|l2 - Duplicate a random table element'''
123    if table_level == 'l1':
124        offset = qed.header['l1_table_offset']
125        table = qed.l1_table
126    elif table_level == 'l2':
127        _, offset = random_table_item(qed.l1_table)
128        table = qed.read_table(offset)
129    else:
130        err('unrecognized sub-command')
131    corrupt_table_duplicate(table)
132    qed.write_table(offset, table)
133
134def cmd_invalidate(qed, table_level):
135    '''invalidate l1|l2 - Plant an invalid table element at random'''
136    if table_level == 'l1':
137        offset = qed.header['l1_table_offset']
138        table = qed.l1_table
139    elif table_level == 'l2':
140        _, offset = random_table_item(qed.l1_table)
141        table = qed.read_table(offset)
142    else:
143        err('unrecognized sub-command')
144    corrupt_table_invalidate(qed, table)
145    qed.write_table(offset, table)
146
147def cmd_need_check(qed, *args):
148    '''need-check [on|off] - Test, set, or clear the QED_F_NEED_CHECK header bit'''
149    if not args:
150        print(bool(qed.header['features'] & QED_F_NEED_CHECK))
151        return
152
153    if args[0] == 'on':
154        qed.header['features'] |= QED_F_NEED_CHECK
155    elif args[0] == 'off':
156        qed.header['features'] &= ~QED_F_NEED_CHECK
157    else:
158        err('unrecognized sub-command')
159    qed.store_header()
160
161def cmd_zero_cluster(qed, pos, *args):
162    '''zero-cluster <pos> [<n>] - Zero data clusters'''
163    pos, n = int(pos), 1
164    if args:
165        if len(args) != 1:
166            err('expected one argument')
167        n = int(args[0])
168
169    for i in xrange(n):
170        l1_index = pos // qed.header['cluster_size'] // len(qed.l1_table)
171        if qed.l1_table[l1_index] == 0:
172            err('no l2 table allocated')
173
174        l2_offset = qed.l1_table[l1_index]
175        l2_table = qed.read_table(l2_offset)
176
177        l2_index = (pos // qed.header['cluster_size']) % len(qed.l1_table)
178        l2_table[l2_index] = 1 # zero the data cluster
179        qed.write_table(l2_offset, l2_table)
180        pos += qed.header['cluster_size']
181
182def cmd_copy_metadata(qed, outfile):
183    '''copy-metadata <outfile> - Copy metadata only (for scrubbing corrupted images)'''
184    out = open(outfile, 'wb')
185
186    # Match file size
187    out.seek(qed.filesize - 1)
188    out.write('\0')
189
190    # Copy header clusters
191    out.seek(0)
192    header_size_bytes = qed.header['header_size'] * qed.header['cluster_size']
193    out.write(qed.raw_pread(0, header_size_bytes))
194
195    # Copy L1 table
196    out.seek(qed.header['l1_table_offset'])
197    s = ''.join(pack_table_elem(x) for x in qed.l1_table)
198    out.write(s)
199
200    # Copy L2 tables
201    for l2_offset in qed.l1_table:
202        if l2_offset == 0:
203            continue
204        l2_table = qed.read_table(l2_offset)
205        out.seek(l2_offset)
206        s = ''.join(pack_table_elem(x) for x in l2_table)
207        out.write(s)
208
209    out.close()
210
211def usage():
212    print('Usage: %s <file> <cmd> [<arg>, ...]' % sys.argv[0])
213    print()
214    print('Supported commands:')
215    for cmd in sorted(x for x in globals() if x.startswith('cmd_')):
216        print(globals()[cmd].__doc__)
217    sys.exit(1)
218
219def main():
220    if len(sys.argv) < 3:
221        usage()
222    filename, cmd = sys.argv[1:3]
223
224    cmd = 'cmd_' + cmd.replace('-', '_')
225    if cmd not in globals():
226        usage()
227
228    qed = QED(open(filename, 'r+b'))
229    try:
230        globals()[cmd](qed, *sys.argv[3:])
231    except TypeError as e:
232        sys.stderr.write(globals()[cmd].__doc__ + '\n')
233        sys.exit(1)
234
235if __name__ == '__main__':
236    main()
237