BmapCreate: update to the latest version
authorArtem Bityutskiy <artem.bityutskiy@intel.com>
Mon, 30 Sep 2013 07:33:15 +0000 (10:33 +0300)
committerGerrit Code Review <gerrit2@otctools.jf.intel.com>
Wed, 16 Oct 2013 07:30:38 +0000 (00:30 -0700)
Use the latest version of BmapCreate. This version generates bmap format 1.4
which is flexible about the checksum algorighm we use. And the default is now
sha256 instead of sha1.

Change-Id: I7aabdab9747d913368a7890be1bde407009b5cca
Signed-off-by: Artem Bityutskiy <artem.bityutskiy@intel.com>
mic/utils/BmapCreate.py

index 65b19a5..7f74bcd 100644 (file)
@@ -1,5 +1,17 @@
-""" This module implements the block map (bmap) creation functionality and
-provides the corresponding API in form of the 'BmapCreate' class.
+# Copyright (c) 2012-2013 Intel, Inc.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License, version 2,
+# as published by the Free Software Foundation.
+#
+# 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.
+
+"""
+This module implements the block map (bmap) creation functionality and provides
+the corresponding API in form of the 'BmapCreate' class.
 
 The idea is that while images files may generally be very large (e.g., 4GiB),
 they may nevertheless contain only little real data, e.g., 512MiB. This data
@@ -22,7 +34,8 @@ etc. The end result should be a sparse file where mapped areas represent useful
 parts of the image and holes represent useless parts of the image, which do not
 have to be copied when copying the image to the target device.
 
-This module uses the FIBMAP ioctl to detect holes. """
+This module uses the FIBMAP ioctl to detect holes.
+"""
 
 # Disable the following pylint recommendations:
 #   *  Too many instance attributes - R0902
@@ -33,8 +46,14 @@ import hashlib
 from mic.utils.misc import human_size
 from mic.utils import Fiemap
 
-# The bmap format version we generate
-SUPPORTED_BMAP_VERSION = "1.3"
+# The bmap format version we generate.
+#
+# Changelog:
+# o 1.3 -> 1.4:
+#   Support SHA256 and SHA512 checksums, in 1.3 only SHA1 was supported.
+#   "BmapFileChecksum" is used instead of "BmapFileSHA1", and "chksum="
+#   attribute is used instead "sha1=". Introduced "ChecksumType" tag.
+SUPPORTED_BMAP_VERSION = "1.4"
 
 _BMAP_START_TEMPLATE = \
 """<?xml version="1.0" ?>
@@ -73,50 +92,36 @@ _BMAP_START_TEMPLATE = \
 """
 
 class Error(Exception):
-    """ A class for exceptions generated by this module. We currently support
-    only one type of exceptions, and we basically throw human-readable problem
-    description in case of errors. """
+    """
+    A class for exceptions generated by this module. We currently support only
+    one type of exceptions, and we basically throw human-readable problem
+    description in case of errors.
+    """
     pass
 
 class BmapCreate:
-    """ This class implements the bmap creation functionality. To generate a
-    bmap for an image (which is supposedly a sparse file), you should first
-    create an instance of 'BmapCreate' and provide:
+    """
+    This class implements the bmap creation functionality. To generate a bmap
+    for an image (which is supposedly a sparse file), you should first create
+    an instance of 'BmapCreate' and provide:
 
     * full path or a file-like object of the image to create bmap for
     * full path or a file object to use for writing the results to
 
     Then you should invoke the 'generate()' method of this class. It will use
-    the FIEMAP ioctl to generate the bmap. """
-
-    def _open_image_file(self):
-        """ Open the image file. """
-
-        try:
-            self._f_image = open(self._image_path, 'rb')
-        except IOError as err:
-            raise Error("cannot open image file '%s': %s" \
-                        % (self._image_path, err))
-
-        self._f_image_needs_close = True
-
-    def _open_bmap_file(self):
-        """ Open the bmap file. """
-
-        try:
-            self._f_bmap = open(self._bmap_path, 'w+')
-        except IOError as err:
-            raise Error("cannot open bmap file '%s': %s" \
-                        % (self._bmap_path, err))
-
-        self._f_bmap_needs_close = True
-
-    def __init__(self, image, bmap):
-        """ Initialize a class instance:
-        * image - full path or a file-like object of the image to create bmap
-                  for
-        * bmap  - full path or a file object to use for writing the resulting
-                  bmap to """
+    the FIEMAP ioctl to generate the bmap.
+    """
+
+    def __init__(self, image, bmap, chksum_type="sha256"):
+        """
+        Initialize a class instance:
+        * image  - full path or a file-like object of the image to create bmap
+                   for
+        * bmap   - full path or a file object to use for writing the resulting
+                   bmap to
+        * chksum - type of the check sum to use in the bmap file (all checksum
+                   types which python's "hashlib" module supports are allowed).
+        """
 
         self.image_size = None
         self.image_size_human = None
@@ -129,11 +134,18 @@ class BmapCreate:
 
         self._mapped_count_pos1 = None
         self._mapped_count_pos2 = None
-        self._sha1_pos = None
+        self._chksum_pos = None
 
         self._f_image_needs_close = False
         self._f_bmap_needs_close = False
 
+        self._cs_type = chksum_type.lower()
+        try:
+            self._cs_len = len(hashlib.new(self._cs_type).hexdigest())
+        except ValueError as err:
+            raise Error("cannot initialize hash function \"%s\": %s" %
+                        (self._cs_type, err))
+
         if hasattr(image, "read"):
             self._f_image = image
             self._image_path = image.name
@@ -153,15 +165,44 @@ class BmapCreate:
         self.image_size = self.fiemap.image_size
         self.image_size_human = human_size(self.image_size)
         if self.image_size == 0:
-            raise Error("cannot generate bmap for zero-sized image file '%s'" \
+            raise Error("cannot generate bmap for zero-sized image file '%s'"
                         % self._image_path)
 
         self.block_size = self.fiemap.block_size
         self.blocks_cnt = self.fiemap.blocks_cnt
 
+    def __del__(self):
+        """The class destructor which closes the opened files."""
+        if self._f_image_needs_close:
+            self._f_image.close()
+        if self._f_bmap_needs_close:
+            self._f_bmap.close()
+
+    def _open_image_file(self):
+        """Open the image file."""
+        try:
+            self._f_image = open(self._image_path, 'rb')
+        except IOError as err:
+            raise Error("cannot open image file '%s': %s"
+                        % (self._image_path, err))
+
+        self._f_image_needs_close = True
+
+    def _open_bmap_file(self):
+        """Open the bmap file."""
+        try:
+            self._f_bmap = open(self._bmap_path, 'w+')
+        except IOError as err:
+            raise Error("cannot open bmap file '%s': %s"
+                        % (self._bmap_path, err))
+
+        self._f_bmap_needs_close = True
+
     def _bmap_file_start(self):
-        """ A helper function which generates the starting contents of the
-        block map file: the header comment, image size, block size, etc. """
+        """
+        A helper function which generates the starting contents of the block
+        map file: the header comment, image size, block size, etc.
+        """
 
         # We do not know the amount of mapped blocks at the moment, so just put
         # whitespaces instead of real numbers. Assume the longest possible
@@ -187,26 +228,31 @@ class BmapCreate:
         xml  = "%s </MappedBlocksCount>\n\n" % mapped_count
 
         # pylint: disable=C0301
+        xml += "    <!-- Type of checksum used in this file -->\n"
+        xml += "    <ChecksumType> %s </ChecksumType>\n\n" % self._cs_type
+
         xml += "    <!-- The checksum of this bmap file. When it is calculated, the value of\n"
-        xml += "         the SHA1 checksum has be zeoro (40 ASCII \"0\" symbols). -->\n"
-        xml += "    <BmapFileSHA1> "
+        xml += "         the checksum has be zero (all ASCII \"0\" symbols).  -->\n"
+        xml += "    <BmapFileChecksum> "
 
         self._f_bmap.write(xml)
-        self._sha1_pos = self._f_bmap.tell()
+        self._chksum_pos = self._f_bmap.tell()
 
-        xml = "0" * 40 + " </BmapFileSHA1>\n\n"
+        xml = "0" * self._cs_len + " </BmapFileChecksum>\n\n"
         xml += "    <!-- The block map which consists of elements which may either be a\n"
-        xml += "         range of blocks or a single block. The 'sha1' attribute (if present)\n"
-        xml += "         is the SHA1 checksum of this blocks range. -->\n"
+        xml += "         range of blocks or a single block. The 'chksum' attribute\n"
+        xml += "         (if present) is the checksum of this blocks range. -->\n"
         xml += "    <BlockMap>\n"
         # pylint: enable=C0301
 
         self._f_bmap.write(xml)
 
     def _bmap_file_end(self):
-        """ A helper function which generates the final parts of the block map
+        """
+        A helper function which generates the final parts of the block map
         file: the ending tags and the information about the amount of mapped
-        blocks. """
+        blocks.
+        """
 
         xml =  "    </BlockMap>\n"
         xml += "</bmap>\n"
@@ -214,26 +260,30 @@ class BmapCreate:
         self._f_bmap.write(xml)
 
         self._f_bmap.seek(self._mapped_count_pos1)
-        self._f_bmap.write("%s or %.1f%%" % \
-                           (self.mapped_size_human, self.mapped_percent))
+        self._f_bmap.write("%s or %.1f%%"
+                           (self.mapped_size_human, self.mapped_percent))
 
         self._f_bmap.seek(self._mapped_count_pos2)
         self._f_bmap.write("%u" % self.mapped_cnt)
 
         self._f_bmap.seek(0)
-        sha1 = hashlib.sha1(self._f_bmap.read()).hexdigest()
-        self._f_bmap.seek(self._sha1_pos)
-        self._f_bmap.write("%s" % sha1)
-
-    def _calculate_sha1(self, first, last):
-        """ A helper function which calculates SHA1 checksum for the range of
-        blocks of the image file: from block 'first' to block 'last'. """
+        hash_obj = hashlib.new(self._cs_type)
+        hash_obj.update(self._f_bmap.read())
+        chksum = hash_obj.hexdigest()
+        self._f_bmap.seek(self._chksum_pos)
+        self._f_bmap.write("%s" % chksum)
+
+    def _calculate_chksum(self, first, last):
+        """
+        A helper function which calculates checksum for the range of blocks of
+        the image file: from block 'first' to block 'last'.
+        """
 
         start = first * self.block_size
         end = (last + 1) * self.block_size
 
         self._f_image.seek(start)
-        hash_obj = hashlib.new("sha1")
+        hash_obj = hashlib.new(self._cs_type)
 
         chunk_size = 1024*1024
         to_read = end - start
@@ -248,9 +298,11 @@ class BmapCreate:
 
         return hash_obj.hexdigest()
 
-    def generate(self, include_checksums = True):
-        """ Generate bmap for the image file. If 'include_checksums' is 'True',
-        also generate SHA1 checksums for block ranges. """
+    def generate(self, include_checksums=True):
+        """
+        Generate bmap for the image file. If 'include_checksums' is 'True',
+        also generate checksums for block ranges.
+        """
 
         # Save image file position in order to restore it at the end
         image_pos = self._f_image.tell()
@@ -263,17 +315,17 @@ class BmapCreate:
         for first, last in self.fiemap.get_mapped_ranges(0, self.blocks_cnt):
             self.mapped_cnt += last - first + 1
             if include_checksums:
-                sha1 = self._calculate_sha1(first, last)
-                sha1 = " sha1=\"%s\"" % sha1
+                chksum = self._calculate_chksum(first, last)
+                chksum = " chksum=\"%s\"" % chksum
             else:
-                sha1 = ""
+                chksum = ""
 
             if first != last:
-                self._f_bmap.write("        <Range%s> %s-%s </Range>\n" \
-                                   % (sha1, first, last))
+                self._f_bmap.write("        <Range%s> %s-%s </Range>\n"
+                                   % (chksum, first, last))
             else:
-                self._f_bmap.write("        <Range%s> %s </Range>\n" \
-                                   % (sha1, first))
+                self._f_bmap.write("        <Range%s> %s </Range>\n"
+                                   % (chksum, first))
 
         self.mapped_size = self.mapped_cnt * self.block_size
         self.mapped_size_human = human_size(self.mapped_size)
@@ -284,15 +336,7 @@ class BmapCreate:
         try:
             self._f_bmap.flush()
         except IOError as err:
-            raise Error("cannot flush the bmap file '%s': %s" \
+            raise Error("cannot flush the bmap file '%s': %s"
                         % (self._bmap_path, err))
 
         self._f_image.seek(image_pos)
-
-    def __del__(self):
-        """ The class destructor which closes the opened files. """
-
-        if self._f_image_needs_close:
-            self._f_image.close()
-        if self._f_bmap_needs_close:
-            self._f_bmap.close()