| # Generator of fuzzed qcow2 images | 
 | # | 
 | # Copyright (C) 2014 Maria Kustova <maria.k@catit.be> | 
 | # | 
 | # This program is free software: you can redistribute it and/or modify | 
 | # it under the terms of the GNU General Public License as published by | 
 | # the Free Software Foundation, either version 2 of the License, or | 
 | # (at your option) any later version. | 
 | # | 
 | # This program is distributed in the hope that it will be useful, | 
 | # but WITHOUT ANY WARRANTY; without even the implied warranty of | 
 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the | 
 | # GNU General Public License for more details. | 
 | # | 
 | # You should have received a copy of the GNU General Public License | 
 | # along with this program.  If not, see <http://www.gnu.org/licenses/>. | 
 | # | 
 |  | 
 | import random | 
 | import struct | 
 | from . import fuzz | 
 | from math import ceil | 
 | from os import urandom | 
 | from itertools import chain | 
 |  | 
 | MAX_IMAGE_SIZE = 10 * (1 << 20) | 
 | # Standard sizes | 
 | UINT32_S = 4 | 
 | UINT64_S = 8 | 
 |  | 
 |  | 
 | class Field(object): | 
 |  | 
 |     """Atomic image element (field). | 
 |  | 
 |     The class represents an image field as quadruple of a data format | 
 |     of value necessary for its packing to binary form, an offset from | 
 |     the beginning of the image, a value and a name. | 
 |  | 
 |     The field can be iterated as a list [format, offset, value, name]. | 
 |     """ | 
 |  | 
 |     __slots__ = ('fmt', 'offset', 'value', 'name') | 
 |  | 
 |     def __init__(self, fmt, offset, val, name): | 
 |         self.fmt = fmt | 
 |         self.offset = offset | 
 |         self.value = val | 
 |         self.name = name | 
 |  | 
 |     def __iter__(self): | 
 |         return iter([self.fmt, self.offset, self.value, self.name]) | 
 |  | 
 |     def __repr__(self): | 
 |         return "Field(fmt=%r, offset=%r, value=%r, name=%r)" % \ | 
 |             (self.fmt, self.offset, self.value, self.name) | 
 |  | 
 |  | 
 | class FieldsList(object): | 
 |  | 
 |     """List of fields. | 
 |  | 
 |     The class allows access to a field in the list by its name. | 
 |     """ | 
 |  | 
 |     def __init__(self, meta_data=None): | 
 |         if meta_data is None: | 
 |             self.data = [] | 
 |         else: | 
 |             self.data = [Field(*f) | 
 |                          for f in meta_data] | 
 |  | 
 |     def __getitem__(self, name): | 
 |         return [x for x in self.data if x.name == name] | 
 |  | 
 |     def __iter__(self): | 
 |         return iter(self.data) | 
 |  | 
 |     def __len__(self): | 
 |         return len(self.data) | 
 |  | 
 |  | 
 | class Image(object): | 
 |  | 
 |     """ Qcow2 image object. | 
 |  | 
 |     This class allows to create qcow2 images with random valid structures and | 
 |     values, fuzz them via external qcow2.fuzz module and write the result to | 
 |     a file. | 
 |     """ | 
 |  | 
 |     def __init__(self, backing_file_name=None): | 
 |         """Create a random valid qcow2 image with the correct header and stored | 
 |         backing file name. | 
 |         """ | 
 |         cluster_bits, self.image_size = self._size_params() | 
 |         self.cluster_size = 1 << cluster_bits | 
 |         self.header = FieldsList() | 
 |         self.backing_file_name = FieldsList() | 
 |         self.backing_file_format = FieldsList() | 
 |         self.feature_name_table = FieldsList() | 
 |         self.end_of_extension_area = FieldsList() | 
 |         self.l2_tables = FieldsList() | 
 |         self.l1_table = FieldsList() | 
 |         self.refcount_table = FieldsList() | 
 |         self.refcount_blocks = FieldsList() | 
 |         self.ext_offset = 0 | 
 |         self.create_header(cluster_bits, backing_file_name) | 
 |         self.set_backing_file_name(backing_file_name) | 
 |         self.data_clusters = self._alloc_data(self.image_size, | 
 |                                               self.cluster_size) | 
 |         # Percentage of fields will be fuzzed | 
 |         self.bias = random.uniform(0.2, 0.5) | 
 |  | 
 |     def __iter__(self): | 
 |         return chain(self.header, self.backing_file_format, | 
 |                      self.feature_name_table, self.end_of_extension_area, | 
 |                      self.backing_file_name, self.l1_table, self.l2_tables, | 
 |                      self.refcount_table, self.refcount_blocks) | 
 |  | 
 |     def create_header(self, cluster_bits, backing_file_name=None): | 
 |         """Generate a random valid header.""" | 
 |         meta_header = [ | 
 |             ['>4s', 0, b"QFI\xfb", 'magic'], | 
 |             ['>I', 4, random.randint(2, 3), 'version'], | 
 |             ['>Q', 8, 0, 'backing_file_offset'], | 
 |             ['>I', 16, 0, 'backing_file_size'], | 
 |             ['>I', 20, cluster_bits, 'cluster_bits'], | 
 |             ['>Q', 24, self.image_size, 'size'], | 
 |             ['>I', 32, 0, 'crypt_method'], | 
 |             ['>I', 36, 0, 'l1_size'], | 
 |             ['>Q', 40, 0, 'l1_table_offset'], | 
 |             ['>Q', 48, 0, 'refcount_table_offset'], | 
 |             ['>I', 56, 0, 'refcount_table_clusters'], | 
 |             ['>I', 60, 0, 'nb_snapshots'], | 
 |             ['>Q', 64, 0, 'snapshots_offset'], | 
 |             ['>Q', 72, 0, 'incompatible_features'], | 
 |             ['>Q', 80, 0, 'compatible_features'], | 
 |             ['>Q', 88, 0, 'autoclear_features'], | 
 |             # Only refcount_order = 4 is supported by current (07.2014) | 
 |             # implementation of QEMU | 
 |             ['>I', 96, 4, 'refcount_order'], | 
 |             ['>I', 100, 0, 'header_length'] | 
 |         ] | 
 |         self.header = FieldsList(meta_header) | 
 |  | 
 |         if self.header['version'][0].value == 2: | 
 |             self.header['header_length'][0].value = 72 | 
 |         else: | 
 |             self.header['incompatible_features'][0].value = \ | 
 |                                                         random.getrandbits(2) | 
 |             self.header['compatible_features'][0].value = random.getrandbits(1) | 
 |             self.header['header_length'][0].value = 104 | 
 |         # Extensions start at the header last field offset and the field size | 
 |         self.ext_offset = struct.calcsize( | 
 |             self.header['header_length'][0].fmt) + \ | 
 |             self.header['header_length'][0].offset | 
 |         end_of_extension_area_len = 2 * UINT32_S | 
 |         free_space = self.cluster_size - self.ext_offset - \ | 
 |                      end_of_extension_area_len | 
 |         # If the backing file name specified and there is enough space for it | 
 |         # in the first cluster, then it's placed in the very end of the first | 
 |         # cluster. | 
 |         if (backing_file_name is not None) and \ | 
 |            (free_space >= len(backing_file_name)): | 
 |             self.header['backing_file_size'][0].value = len(backing_file_name) | 
 |             self.header['backing_file_offset'][0].value = \ | 
 |                                     self.cluster_size - len(backing_file_name) | 
 |  | 
 |     def set_backing_file_name(self, backing_file_name=None): | 
 |         """Add the name of the backing file at the offset specified | 
 |         in the header. | 
 |         """ | 
 |         if (backing_file_name is not None) and \ | 
 |            (not self.header['backing_file_offset'][0].value == 0): | 
 |             data_len = len(backing_file_name) | 
 |             data_fmt = '>' + str(data_len) + 's' | 
 |             self.backing_file_name = FieldsList([ | 
 |                 [data_fmt, self.header['backing_file_offset'][0].value, | 
 |                  backing_file_name, 'bf_name'] | 
 |             ]) | 
 |  | 
 |     def set_backing_file_format(self, backing_file_fmt=None): | 
 |         """Generate the header extension for the backing file format.""" | 
 |         if backing_file_fmt is not None: | 
 |             # Calculation of the free space available in the first cluster | 
 |             end_of_extension_area_len = 2 * UINT32_S | 
 |             high_border = (self.header['backing_file_offset'][0].value or | 
 |                            (self.cluster_size - 1)) - \ | 
 |                 end_of_extension_area_len | 
 |             free_space = high_border - self.ext_offset | 
 |             ext_size = 2 * UINT32_S + ((len(backing_file_fmt) + 7) & ~7) | 
 |  | 
 |             if free_space >= ext_size: | 
 |                 ext_data_len = len(backing_file_fmt) | 
 |                 ext_data_fmt = '>' + str(ext_data_len) + 's' | 
 |                 ext_padding_len = 7 - (ext_data_len - 1) % 8 | 
 |                 self.backing_file_format = FieldsList([ | 
 |                     ['>I', self.ext_offset, 0xE2792ACA, 'ext_magic'], | 
 |                     ['>I', self.ext_offset + UINT32_S, ext_data_len, | 
 |                      'ext_length'], | 
 |                     [ext_data_fmt, self.ext_offset + UINT32_S * 2, | 
 |                      backing_file_fmt, 'bf_format'] | 
 |                 ]) | 
 |                 self.ext_offset = \ | 
 |                         struct.calcsize( | 
 |                             self.backing_file_format['bf_format'][0].fmt) + \ | 
 |                         ext_padding_len + \ | 
 |                         self.backing_file_format['bf_format'][0].offset | 
 |  | 
 |     def create_feature_name_table(self): | 
 |         """Generate a random header extension for names of features used in | 
 |         the image. | 
 |         """ | 
 |         def gen_feat_ids(): | 
 |             """Return random feature type and feature bit.""" | 
 |             return (random.randint(0, 2), random.randint(0, 63)) | 
 |  | 
 |         end_of_extension_area_len = 2 * UINT32_S | 
 |         high_border = (self.header['backing_file_offset'][0].value or | 
 |                        (self.cluster_size - 1)) - \ | 
 |             end_of_extension_area_len | 
 |         free_space = high_border - self.ext_offset | 
 |         # Sum of sizes of 'magic' and 'length' header extension fields | 
 |         ext_header_len = 2 * UINT32_S | 
 |         fnt_entry_size = 6 * UINT64_S | 
 |         num_fnt_entries = min(10, (free_space - ext_header_len) / | 
 |                               fnt_entry_size) | 
 |         if not num_fnt_entries == 0: | 
 |             feature_tables = [] | 
 |             feature_ids = [] | 
 |             inner_offset = self.ext_offset + ext_header_len | 
 |             feat_name = b'some cool feature' | 
 |             while len(feature_tables) < num_fnt_entries * 3: | 
 |                 feat_type, feat_bit = gen_feat_ids() | 
 |                 # Remove duplicates | 
 |                 while (feat_type, feat_bit) in feature_ids: | 
 |                     feat_type, feat_bit = gen_feat_ids() | 
 |                 feature_ids.append((feat_type, feat_bit)) | 
 |                 feat_fmt = '>' + str(len(feat_name)) + 's' | 
 |                 feature_tables += [['B', inner_offset, | 
 |                                     feat_type, 'feature_type'], | 
 |                                    ['B', inner_offset + 1, feat_bit, | 
 |                                     'feature_bit_number'], | 
 |                                    [feat_fmt, inner_offset + 2, | 
 |                                     feat_name, 'feature_name'] | 
 |                 ] | 
 |                 inner_offset += fnt_entry_size | 
 |             # No padding for the extension is necessary, because | 
 |             # the extension length is multiple of 8 | 
 |             self.feature_name_table = FieldsList([ | 
 |                 ['>I', self.ext_offset, 0x6803f857, 'ext_magic'], | 
 |                 # One feature table contains 3 fields and takes 48 bytes | 
 |                 ['>I', self.ext_offset + UINT32_S, | 
 |                  len(feature_tables) // 3 * 48, 'ext_length'] | 
 |             ] + feature_tables) | 
 |             self.ext_offset = inner_offset | 
 |  | 
 |     def set_end_of_extension_area(self): | 
 |         """Generate a mandatory header extension marking end of header | 
 |         extensions. | 
 |         """ | 
 |         self.end_of_extension_area = FieldsList([ | 
 |             ['>I', self.ext_offset, 0, 'ext_magic'], | 
 |             ['>I', self.ext_offset + UINT32_S, 0, 'ext_length'] | 
 |         ]) | 
 |  | 
 |     def create_l_structures(self): | 
 |         """Generate random valid L1 and L2 tables.""" | 
 |         def create_l2_entry(host, guest, l2_cluster): | 
 |             """Generate one L2 entry.""" | 
 |             offset = l2_cluster * self.cluster_size | 
 |             l2_size = self.cluster_size // UINT64_S | 
 |             entry_offset = offset + UINT64_S * (guest % l2_size) | 
 |             cluster_descriptor = host * self.cluster_size | 
 |             if not self.header['version'][0].value == 2: | 
 |                 cluster_descriptor += random.randint(0, 1) | 
 |             # While snapshots are not supported, bit #63 = 1 | 
 |             # Compressed clusters are not supported => bit #62 = 0 | 
 |             entry_val = (1 << 63) + cluster_descriptor | 
 |             return ['>Q', entry_offset, entry_val, 'l2_entry'] | 
 |  | 
 |         def create_l1_entry(l2_cluster, l1_offset, guest): | 
 |             """Generate one L1 entry.""" | 
 |             l2_size = self.cluster_size // UINT64_S | 
 |             entry_offset = l1_offset + UINT64_S * (guest // l2_size) | 
 |             # While snapshots are not supported bit #63 = 1 | 
 |             entry_val = (1 << 63) + l2_cluster * self.cluster_size | 
 |             return ['>Q', entry_offset, entry_val, 'l1_entry'] | 
 |  | 
 |         if len(self.data_clusters) == 0: | 
 |             # All metadata for an empty guest image needs 4 clusters: | 
 |             # header, rfc table, rfc block, L1 table. | 
 |             # Header takes cluster #0, other clusters ##1-3 can be used | 
 |             l1_offset = random.randint(1, 3) * self.cluster_size | 
 |             l1 = [['>Q', l1_offset, 0, 'l1_entry']] | 
 |             l2 = [] | 
 |         else: | 
 |             meta_data = self._get_metadata() | 
 |             guest_clusters = random.sample(range(self.image_size // | 
 |                                                  self.cluster_size), | 
 |                                            len(self.data_clusters)) | 
 |             # Number of entries in a L1/L2 table | 
 |             l_size = self.cluster_size // UINT64_S | 
 |             # Number of clusters necessary for L1 table | 
 |             l1_size = int(ceil((max(guest_clusters) + 1) / float(l_size**2))) | 
 |             l1_start = self._get_adjacent_clusters(self.data_clusters | | 
 |                                                    meta_data, l1_size) | 
 |             meta_data |= set(range(l1_start, l1_start + l1_size)) | 
 |             l1_offset = l1_start * self.cluster_size | 
 |             # Indices of L2 tables | 
 |             l2_ids = [] | 
 |             # Host clusters allocated for L2 tables | 
 |             l2_clusters = [] | 
 |             # L1 entries | 
 |             l1 = [] | 
 |             # L2 entries | 
 |             l2 = [] | 
 |             for host, guest in zip(self.data_clusters, guest_clusters): | 
 |                 l2_id = guest // l_size | 
 |                 if l2_id not in l2_ids: | 
 |                     l2_ids.append(l2_id) | 
 |                     l2_clusters.append(self._get_adjacent_clusters( | 
 |                         self.data_clusters | meta_data | set(l2_clusters), | 
 |                         1)) | 
 |                     l1.append(create_l1_entry(l2_clusters[-1], l1_offset, | 
 |                                               guest)) | 
 |                 l2.append(create_l2_entry(host, guest, | 
 |                                           l2_clusters[l2_ids.index(l2_id)])) | 
 |         self.l2_tables = FieldsList(l2) | 
 |         self.l1_table = FieldsList(l1) | 
 |         self.header['l1_size'][0].value = int(ceil(UINT64_S * self.image_size / | 
 |                                                 float(self.cluster_size**2))) | 
 |         self.header['l1_table_offset'][0].value = l1_offset | 
 |  | 
 |     def create_refcount_structures(self): | 
 |         """Generate random refcount blocks and refcount table.""" | 
 |         def allocate_rfc_blocks(data, size): | 
 |             """Return indices of clusters allocated for refcount blocks.""" | 
 |             cluster_ids = set() | 
 |             diff = block_ids = set([x // size for x in data]) | 
 |             while len(diff) != 0: | 
 |                 # Allocate all yet not allocated clusters | 
 |                 new = self._get_available_clusters(data | cluster_ids, | 
 |                                                    len(diff)) | 
 |                 # Indices of new refcount blocks necessary to cover clusters | 
 |                 # in 'new' | 
 |                 diff = set([x // size for x in new]) - block_ids | 
 |                 cluster_ids |= new | 
 |                 block_ids |= diff | 
 |             return cluster_ids, block_ids | 
 |  | 
 |         def allocate_rfc_table(data, init_blocks, block_size): | 
 |             """Return indices of clusters allocated for the refcount table | 
 |             and updated indices of clusters allocated for blocks and indices | 
 |             of blocks. | 
 |             """ | 
 |             blocks = set(init_blocks) | 
 |             clusters = set() | 
 |             # Number of entries in one cluster of the refcount table | 
 |             size = self.cluster_size // UINT64_S | 
 |             # Number of clusters necessary for the refcount table based on | 
 |             # the current number of refcount blocks | 
 |             table_size = int(ceil((max(blocks) + 1) / float(size))) | 
 |             # Index of the first cluster of the refcount table | 
 |             table_start = self._get_adjacent_clusters(data, table_size + 1) | 
 |             # Clusters allocated for the current length of the refcount table | 
 |             table_clusters = set(range(table_start, table_start + table_size)) | 
 |             # Clusters allocated for the refcount table including | 
 |             # last optional one for potential l1 growth | 
 |             table_clusters_allocated = set(range(table_start, table_start + | 
 |                                                  table_size + 1)) | 
 |             # New refcount blocks necessary for clusters occupied by the | 
 |             # refcount table | 
 |             diff = set([c // block_size for c in table_clusters]) - blocks | 
 |             blocks |= diff | 
 |             while len(diff) != 0: | 
 |                 # Allocate clusters for new refcount blocks | 
 |                 new = self._get_available_clusters((data | clusters) | | 
 |                                                    table_clusters_allocated, | 
 |                                                    len(diff)) | 
 |                 # Indices of new refcount blocks necessary to cover | 
 |                 # clusters in 'new' | 
 |                 diff = set([x // block_size for x in new]) - blocks | 
 |                 clusters |= new | 
 |                 blocks |= diff | 
 |                 # Check if the refcount table needs one more cluster | 
 |                 if int(ceil((max(blocks) + 1) / float(size))) > table_size: | 
 |                     new_block_id = (table_start + table_size) // block_size | 
 |                     # Check if the additional table cluster needs | 
 |                     # one more refcount block | 
 |                     if new_block_id not in blocks: | 
 |                         diff.add(new_block_id) | 
 |                     table_clusters.add(table_start + table_size) | 
 |                     table_size += 1 | 
 |             return table_clusters, blocks, clusters | 
 |  | 
 |         def create_table_entry(table_offset, block_cluster, block_size, | 
 |                                cluster): | 
 |             """Generate a refcount table entry.""" | 
 |             offset = table_offset + UINT64_S * (cluster // block_size) | 
 |             return ['>Q', offset, block_cluster * self.cluster_size, | 
 |                     'refcount_table_entry'] | 
 |  | 
 |         def create_block_entry(block_cluster, block_size, cluster): | 
 |             """Generate a list of entries for the current block.""" | 
 |             entry_size = self.cluster_size // block_size | 
 |             offset = block_cluster * self.cluster_size | 
 |             entry_offset = offset + entry_size * (cluster % block_size) | 
 |             # While snapshots are not supported all refcounts are set to 1 | 
 |             return ['>H', entry_offset, 1, 'refcount_block_entry'] | 
 |         # Size of a block entry in bits | 
 |         refcount_bits = 1 << self.header['refcount_order'][0].value | 
 |         # Number of refcount entries per refcount block | 
 |         # Convert self.cluster_size from bytes to bits to have the same | 
 |         # base for the numerator and denominator | 
 |         block_size = self.cluster_size * 8 // refcount_bits | 
 |         meta_data = self._get_metadata() | 
 |         if len(self.data_clusters) == 0: | 
 |             # All metadata for an empty guest image needs 4 clusters: | 
 |             # header, rfc table, rfc block, L1 table. | 
 |             # Header takes cluster #0, other clusters ##1-3 can be used | 
 |             block_clusters = set([random.choice(list(set(range(1, 4)) - | 
 |                                                      meta_data))]) | 
 |             block_ids = set([0]) | 
 |             table_clusters = set([random.choice(list(set(range(1, 4)) - | 
 |                                                      meta_data - | 
 |                                                      block_clusters))]) | 
 |         else: | 
 |             block_clusters, block_ids = \ | 
 |                                 allocate_rfc_blocks(self.data_clusters | | 
 |                                                     meta_data, block_size) | 
 |             table_clusters, block_ids, new_clusters = \ | 
 |                                     allocate_rfc_table(self.data_clusters | | 
 |                                                        meta_data | | 
 |                                                        block_clusters, | 
 |                                                        block_ids, | 
 |                                                        block_size) | 
 |             block_clusters |= new_clusters | 
 |  | 
 |         meta_data |= block_clusters | table_clusters | 
 |         table_offset = min(table_clusters) * self.cluster_size | 
 |         block_id = None | 
 |         # Clusters allocated for refcount blocks | 
 |         block_clusters = list(block_clusters) | 
 |         # Indices of refcount blocks | 
 |         block_ids = list(block_ids) | 
 |         # Refcount table entries | 
 |         rfc_table = [] | 
 |         # Refcount entries | 
 |         rfc_blocks = [] | 
 |  | 
 |         for cluster in sorted(self.data_clusters | meta_data): | 
 |             if cluster // block_size != block_id: | 
 |                 block_id = cluster // block_size | 
 |                 block_cluster = block_clusters[block_ids.index(block_id)] | 
 |                 rfc_table.append(create_table_entry(table_offset, | 
 |                                                     block_cluster, | 
 |                                                     block_size, cluster)) | 
 |             rfc_blocks.append(create_block_entry(block_cluster, block_size, | 
 |                                                  cluster)) | 
 |         self.refcount_table = FieldsList(rfc_table) | 
 |         self.refcount_blocks = FieldsList(rfc_blocks) | 
 |  | 
 |         self.header['refcount_table_offset'][0].value = table_offset | 
 |         self.header['refcount_table_clusters'][0].value = len(table_clusters) | 
 |  | 
 |     def fuzz(self, fields_to_fuzz=None): | 
 |         """Fuzz an image by corrupting values of a random subset of its fields. | 
 |  | 
 |         Without parameters the method fuzzes an entire image. | 
 |  | 
 |         If 'fields_to_fuzz' is specified then only fields in this list will be | 
 |         fuzzed. 'fields_to_fuzz' can contain both individual fields and more | 
 |         general image elements as a header or tables. | 
 |  | 
 |         In the first case the field will be fuzzed always. | 
 |         In the second a random subset of fields will be selected and fuzzed. | 
 |         """ | 
 |         def coin(): | 
 |             """Return boolean value proportional to a portion of fields to be | 
 |             fuzzed. | 
 |             """ | 
 |             return random.random() < self.bias | 
 |  | 
 |         if fields_to_fuzz is None: | 
 |             for field in self: | 
 |                 if coin(): | 
 |                     field.value = getattr(fuzz, field.name)(field.value) | 
 |         else: | 
 |             for item in fields_to_fuzz: | 
 |                 if len(item) == 1: | 
 |                     for field in getattr(self, item[0]): | 
 |                         if coin(): | 
 |                             field.value = getattr(fuzz, | 
 |                                                   field.name)(field.value) | 
 |                 else: | 
 |                     # If fields with the requested name were not generated | 
 |                     # getattr(self, item[0])[item[1]] returns an empty list | 
 |                     for field in getattr(self, item[0])[item[1]]: | 
 |                         field.value = getattr(fuzz, field.name)(field.value) | 
 |  | 
 |     def write(self, filename): | 
 |         """Write an entire image to the file.""" | 
 |         image_file = open(filename, 'wb') | 
 |         for field in self: | 
 |             image_file.seek(field.offset) | 
 |             image_file.write(struct.pack(field.fmt, field.value)) | 
 |  | 
 |         for cluster in sorted(self.data_clusters): | 
 |             image_file.seek(cluster * self.cluster_size) | 
 |             image_file.write(urandom(self.cluster_size)) | 
 |  | 
 |         # Align the real image size to the cluster size | 
 |         image_file.seek(0, 2) | 
 |         size = image_file.tell() | 
 |         rounded = (size + self.cluster_size - 1) & ~(self.cluster_size - 1) | 
 |         if rounded > size: | 
 |             image_file.seek(rounded - 1) | 
 |             image_file.write(b'\x00') | 
 |         image_file.close() | 
 |  | 
 |     @staticmethod | 
 |     def _size_params(): | 
 |         """Generate a random image size aligned to a random correct | 
 |         cluster size. | 
 |         """ | 
 |         cluster_bits = random.randrange(9, 21) | 
 |         cluster_size = 1 << cluster_bits | 
 |         img_size = random.randrange(0, MAX_IMAGE_SIZE + 1, cluster_size) | 
 |         return (cluster_bits, img_size) | 
 |  | 
 |     @staticmethod | 
 |     def _get_available_clusters(used, number): | 
 |         """Return a set of indices of not allocated clusters. | 
 |  | 
 |         'used' contains indices of currently allocated clusters. | 
 |         All clusters that cannot be allocated between 'used' clusters will have | 
 |         indices appended to the end of 'used'. | 
 |         """ | 
 |         append_id = max(used) + 1 | 
 |         free = set(range(1, append_id)) - used | 
 |         if len(free) >= number: | 
 |             return set(random.sample(free, number)) | 
 |         else: | 
 |             return free | set(range(append_id, append_id + number - len(free))) | 
 |  | 
 |     @staticmethod | 
 |     def _get_adjacent_clusters(used, size): | 
 |         """Return an index of the first cluster in the sequence of free ones. | 
 |  | 
 |         'used' contains indices of currently allocated clusters. 'size' is the | 
 |         length of the sequence of free clusters. | 
 |         If the sequence of 'size' is not available between 'used' clusters, its | 
 |         first index will be append to the end of 'used'. | 
 |         """ | 
 |         def get_cluster_id(lst, length): | 
 |             """Return the first index of the sequence of the specified length | 
 |             or None if the sequence cannot be inserted in the list. | 
 |             """ | 
 |             if len(lst) != 0: | 
 |                 pairs = [] | 
 |                 pair = (lst[0], 1) | 
 |                 for i in range(1, len(lst)): | 
 |                     if lst[i] == lst[i-1] + 1: | 
 |                         pair = (lst[i], pair[1] + 1) | 
 |                     else: | 
 |                         pairs.append(pair) | 
 |                         pair = (lst[i], 1) | 
 |                 pairs.append(pair) | 
 |                 random.shuffle(pairs) | 
 |                 for x, s in pairs: | 
 |                     if s >= length: | 
 |                         return x - length + 1 | 
 |             return None | 
 |  | 
 |         append_id = max(used) + 1 | 
 |         free = list(set(range(1, append_id)) - used) | 
 |         idx = get_cluster_id(free, size) | 
 |         if idx is None: | 
 |             return append_id | 
 |         else: | 
 |             return idx | 
 |  | 
 |     @staticmethod | 
 |     def _alloc_data(img_size, cluster_size): | 
 |         """Return a set of random indices of clusters allocated for guest data. | 
 |         """ | 
 |         num_of_cls = img_size // cluster_size | 
 |         return set(random.sample(range(1, num_of_cls + 1), | 
 |                                  random.randint(0, num_of_cls))) | 
 |  | 
 |     def _get_metadata(self): | 
 |         """Return indices of clusters allocated for image metadata.""" | 
 |         ids = set() | 
 |         for x in self: | 
 |             ids.add(x.offset // self.cluster_size) | 
 |         return ids | 
 |  | 
 |  | 
 | def create_image(test_img_path, backing_file_name=None, backing_file_fmt=None, | 
 |                  fields_to_fuzz=None): | 
 |     """Create a fuzzed image and write it to the specified file.""" | 
 |     image = Image(backing_file_name.encode()) | 
 |     image.set_backing_file_format(backing_file_fmt.encode()) | 
 |     image.create_feature_name_table() | 
 |     image.set_end_of_extension_area() | 
 |     image.create_l_structures() | 
 |     image.create_refcount_structures() | 
 |     image.fuzz(fields_to_fuzz) | 
 |     image.write(test_img_path) | 
 |     return image.image_size |