aboutsummaryrefslogtreecommitdiff
path: root/tests/qemu-iotests/qed.py
blob: d6bec9606961687d5356bbff05d36bf38ef09e95 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
#!/usr/bin/env python3
#
# Tool to manipulate QED image files
#
# Copyright (C) 2010 IBM, Corp.
#
# Authors:
#  Stefan Hajnoczi <stefanha@linux.vnet.ibm.com>
#
# This work is licensed under the terms of the GNU GPL, version 2 or later.
# See the COPYING file in the top-level directory.

import sys
import struct
import random
import optparse

# This can be used as a module
__all__ = ['QED_F_NEED_CHECK', 'QED']

QED_F_NEED_CHECK = 0x02

header_fmt = '<IIIIQQQQQII'
header_size = struct.calcsize(header_fmt)
field_names = ['magic', 'cluster_size', 'table_size',
               'header_size', 'features', 'compat_features',
               'autoclear_features', 'l1_table_offset', 'image_size',
               'backing_filename_offset', 'backing_filename_size']
table_elem_fmt = '<Q'
table_elem_size = struct.calcsize(table_elem_fmt)

def err(msg):
    sys.stderr.write(msg + '\n')
    sys.exit(1)

def unpack_header(s):
    fields = struct.unpack(header_fmt, s)
    return dict((field_names[idx], val) for idx, val in enumerate(fields))

def pack_header(header):
    fields = tuple(header[x] for x in field_names)
    return struct.pack(header_fmt, *fields)

def unpack_table_elem(s):
    return struct.unpack(table_elem_fmt, s)[0]

def pack_table_elem(elem):
    return struct.pack(table_elem_fmt, elem)

class QED(object):
    def __init__(self, f):
        self.f = f

        self.f.seek(0, 2)
        self.filesize = f.tell()

        self.load_header()
        self.load_l1_table()

    def raw_pread(self, offset, size):
        self.f.seek(offset)
        return self.f.read(size)

    def raw_pwrite(self, offset, data):
        self.f.seek(offset)
        return self.f.write(data)

    def load_header(self):
        self.header = unpack_header(self.raw_pread(0, header_size))

    def store_header(self):
        self.raw_pwrite(0, pack_header(self.header))

    def read_table(self, offset):
        size = self.header['table_size'] * self.header['cluster_size']
        s = self.raw_pread(offset, size)
        table = [unpack_table_elem(s[i:i + table_elem_size]) for i in xrange(0, size, table_elem_size)]
        return table

    def load_l1_table(self):
        self.l1_table = self.read_table(self.header['l1_table_offset'])
        self.table_nelems = self.header['table_size'] * self.header['cluster_size'] // table_elem_size

    def write_table(self, offset, table):
        s = ''.join(pack_table_elem(x) for x in table)
        self.raw_pwrite(offset, s)

def random_table_item(table):
    vals = [(index, offset) for index, offset in enumerate(table) if offset != 0]
    if not vals:
        err('cannot pick random item because table is empty')
    return random.choice(vals)

def corrupt_table_duplicate(table):
    '''Corrupt a table by introducing a duplicate offset'''
    victim_idx, victim_val = random_table_item(table)
    unique_vals = set(table)
    if len(unique_vals) == 1:
        err('no duplication corruption possible in table')
    dup_val = random.choice(list(unique_vals.difference([victim_val])))
    table[victim_idx] = dup_val

def corrupt_table_invalidate(qed, table):
    '''Corrupt a table by introducing an invalid offset'''
    index, _ = random_table_item(table)
    table[index] = qed.filesize + random.randint(0, 100 * 1024 * 1024 * 1024 * 1024)

def cmd_show(qed, *args):
    '''show [header|l1|l2 <offset>]- Show header or l1/l2 tables'''
    if not args or args[0] == 'header':
        print(qed.header)
    elif args[0] == 'l1':
        print(qed.l1_table)
    elif len(args) == 2 and args[0] == 'l2':
        offset = int(args[1])
        print(qed.read_table(offset))
    else:
        err('unrecognized sub-command')

def cmd_duplicate(qed, table_level):
    '''duplicate l1|l2 - Duplicate a random table element'''
    if table_level == 'l1':
        offset = qed.header['l1_table_offset']
        table = qed.l1_table
    elif table_level == 'l2':
        _, offset = random_table_item(qed.l1_table)
        table = qed.read_table(offset)
    else:
        err('unrecognized sub-command')
    corrupt_table_duplicate(table)
    qed.write_table(offset, table)

def cmd_invalidate(qed, table_level):
    '''invalidate l1|l2 - Plant an invalid table element at random'''
    if table_level == 'l1':
        offset = qed.header['l1_table_offset']
        table = qed.l1_table
    elif table_level == 'l2':
        _, offset = random_table_item(qed.l1_table)
        table = qed.read_table(offset)
    else:
        err('unrecognized sub-command')
    corrupt_table_invalidate(qed, table)
    qed.write_table(offset, table)

def cmd_need_check(qed, *args):
    '''need-check [on|off] - Test, set, or clear the QED_F_NEED_CHECK header bit'''
    if not args:
        print(bool(qed.header['features'] & QED_F_NEED_CHECK))
        return

    if args[0] == 'on':
        qed.header['features'] |= QED_F_NEED_CHECK
    elif args[0] == 'off':
        qed.header['features'] &= ~QED_F_NEED_CHECK
    else:
        err('unrecognized sub-command')
    qed.store_header()

def cmd_zero_cluster(qed, pos, *args):
    '''zero-cluster <pos> [<n>] - Zero data clusters'''
    pos, n = int(pos), 1
    if args:
        if len(args) != 1:
            err('expected one argument')
        n = int(args[0])

    for i in xrange(n):
        l1_index = pos // qed.header['cluster_size'] // len(qed.l1_table)
        if qed.l1_table[l1_index] == 0:
            err('no l2 table allocated')

        l2_offset = qed.l1_table[l1_index]
        l2_table = qed.read_table(l2_offset)

        l2_index = (pos // qed.header['cluster_size']) % len(qed.l1_table)
        l2_table[l2_index] = 1 # zero the data cluster
        qed.write_table(l2_offset, l2_table)
        pos += qed.header['cluster_size']

def cmd_copy_metadata(qed, outfile):
    '''copy-metadata <outfile> - Copy metadata only (for scrubbing corrupted images)'''
    out = open(outfile, 'wb')

    # Match file size
    out.seek(qed.filesize - 1)
    out.write('\0')

    # Copy header clusters
    out.seek(0)
    header_size_bytes = qed.header['header_size'] * qed.header['cluster_size']
    out.write(qed.raw_pread(0, header_size_bytes))

    # Copy L1 table
    out.seek(qed.header['l1_table_offset'])
    s = ''.join(pack_table_elem(x) for x in qed.l1_table)
    out.write(s)

    # Copy L2 tables
    for l2_offset in qed.l1_table:
        if l2_offset == 0:
            continue
        l2_table = qed.read_table(l2_offset)
        out.seek(l2_offset)
        s = ''.join(pack_table_elem(x) for x in l2_table)
        out.write(s)

    out.close()

def usage():
    print('Usage: %s <file> <cmd> [<arg>, ...]' % sys.argv[0])
    print()
    print('Supported commands:')
    for cmd in sorted(x for x in globals() if x.startswith('cmd_')):
        print(globals()[cmd].__doc__)
    sys.exit(1)

def main():
    if len(sys.argv) < 3:
        usage()
    filename, cmd = sys.argv[1:3]

    cmd = 'cmd_' + cmd.replace('-', '_')
    if cmd not in globals():
        usage()

    qed = QED(open(filename, 'r+b'))
    try:
        globals()[cmd](qed, *sys.argv[3:])
    except TypeError as e:
        sys.stderr.write(globals()[cmd].__doc__ + '\n')
        sys.exit(1)

if __name__ == '__main__':
    main()