Add an Image method for computing the MD5 checksum

This used to be a class in the util module, but since this is only used
on the Image, it suits better there.

image_creator/dialog_menu.py

@@ -27,7 +27,7 @@ import re
 import time
 
 from image_creator import __version__ as version
-from image_creator.util import MD5, FatalError, virtio_versions
+from image_creator.util import FatalError, virtio_versions
 from image_creator.output.dialog import GaugeOutput, InfoBoxOutput
 from image_creator.kamaki_wrapper import Kamaki, ClientError
 from image_creator.help import get_help_file
@@ -149,8 +149,7 @@ def upload_image(session):
         kamaki.out = out
         try:
             if 'checksum' not in session:
-                md5 = MD5(out)
-                session['checksum'] = md5.compute(image.device, image.size)
+                session['checksum'] = image.md5()
 
             try:
                 # Upload image file

image_creator/dialog_util.py

@@ -24,7 +24,6 @@ import stat
 import re
 import json
 from image_creator.output.dialog import GaugeOutput
-from image_creator.util import MD5
 from image_creator.kamaki_wrapper import Kamaki
 
 SMALL_WIDTH = 60
@@ -198,8 +197,7 @@ def extract_image(session):
             out.add(gauge)
             try:
                 if "checksum" not in session:
-                    md5 = MD5(out)
-                    session['checksum'] = md5.compute(image.device, image.size)
+                    session['checksum'] = image.md5()
 
                 # Extract image file
                 image.dump(path)

image_creator/dialog_wizard.py

@@ -25,7 +25,7 @@ import json
 import re
 
 from image_creator.kamaki_wrapper import Kamaki, ClientError
-from image_creator.util import MD5, FatalError, virtio_versions
+from image_creator.util import FatalError, virtio_versions
 from image_creator.output.cli import OutputWthProgress
 from image_creator.dialog_util import extract_image, update_background_title, \
     add_cloud, edit_cloud, update_sysprep_param
@@ -461,7 +461,7 @@ def create_image(session, answers):
         metadata['DESCRIPTION'] = answers['ImageDescription']
 
         # MD5
-        session['checksum'] = MD5(image.out).compute(image.device, image.size)
+        session['checksum'] = image.md5()
 
         image.out.output()
         try:

image_creator/image.py

@@ -21,6 +21,7 @@ from image_creator.os_type import os_cls
 
 import re
 import guestfs
+import hashlib
 from sendfile import sendfile
 
 
@@ -391,14 +392,13 @@ class Image(object):
         partition table. Empty space in the end of the device will be ignored.
         """
         MB = 2 ** 20
-        blocksize = 4 * MB  # 4MB
-        size = self.size
-        progr_size = (size + MB - 1) // MB  # in MB
+        blocksize = 2 ** 22  # 4MB
+        progr_size = (self.size + MB - 1) // MB  # in MB
         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 
-        with open(self.device, 'r') as src:
-            with open(outfile, "w") as dst:
-                left = size
+        with open(self.device, 'rb') as src:
+            with open(outfile, "wb") as dst:
+                left = self.size
                 offset = 0
                 progressbar.next()
                 while left > 0:
@@ -414,7 +414,31 @@ class Image(object):
 
                     offset += sent
                     left -= sent
-                    progressbar.goto((size - left) // MB)
+                    progressbar.goto((self.size - left) // MB)
         progressbar.success('image file %s was successfully created' % outfile)
 
+    def md5(self):
+        """Computes the MD5 checksum of the image"""
+
+        MB = 2 ** 20
+        blocksize = 2 ** 22  # 4MB
+        progr_size = ((self.size + MB - 1) // MB)  # in MB
+        progressbar = self.out.Progress(progr_size, "Calculating md5sum", 'mb')
+        md5 = hashlib.md5()
+
+        with open(self.device, "rb") as src:
+            left = self.size
+            while left > 0:
+                length = min(left, blocksize)
+                data = src.read(length)
+                md5.update(data)
+                left -= length
+                progressbar.goto((self.size - left) // MB)
+
+        checksum = md5.hexdigest()
+        progressbar.success(checksum)
+
+        return checksum
+
+
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/main.py

@@ -22,7 +22,7 @@ snf-image-creator program.
 
 from image_creator import __version__ as version
 from image_creator.disk import Disk
-from image_creator.util import FatalError, MD5
+from image_creator.util import FatalError
 from image_creator.output.cli import SilentOutput, SimpleOutput, \
     OutputWthProgress
 from image_creator.kamaki_wrapper import Kamaki, ClientError
@@ -314,8 +314,7 @@ def image_creator():
         # Add command line metadata to the collected ones...
         metadata.update(options.metadata)
 
-        md5 = MD5(out)
-        checksum = md5.compute(image.device, image.size)
+        checksum = image.md5()
 
         metastring = unicode(json.dumps(
             {'properties': metadata,

image_creator/util.py

@@ -20,7 +20,6 @@ the package.
 """
 
 import sh
-import hashlib
 import time
 import os
 import re
@@ -97,35 +96,6 @@ def free_space(dirname):
     return stat.f_bavail * stat.f_frsize
 
 
-class MD5:
-    """Represents MD5 computations"""
-    def __init__(self, output):
-        """Create an MD5 instance"""
-        self.out = output
-
-    def compute(self, filename, size):
-        """Compute the MD5 checksum of a file"""
-        MB = 2 ** 20
-        BLOCKSIZE = 4 * MB  # 4MB
-
-        prog_size = ((size + MB - 1) // MB)  # in MB
-        progressbar = self.out.Progress(prog_size, "Calculating md5sum", 'mb')
-        md5 = hashlib.md5()
-        with open(filename, "r") as src:
-            left = size
-            while left > 0:
-                length = min(left, BLOCKSIZE)
-                data = src.read(length)
-                md5.update(data)
-                left -= length
-                progressbar.goto((size - left) // MB)
-
-        checksum = md5.hexdigest()
-        progressbar.success(checksum)
-
-        return checksum
-
-
 def virtio_versions(virtio_state):
     """Returns the versions of the drivers defined by the virtio state"""