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