Utils: Create a generic hash function for autotest

index bda72e2..a095e06 100644
@@ -4,6 +4,10 @@
 import os, pickle, random, re, resource, select, shutil, signal, StringIO
 import socket, struct, subprocess, sys, time, textwrap, urlparse
 import warnings, smtplib, logging, urllib2
+try:
+    import hashlib
+except ImportError:
+    import md5, sha
 from autotest_lib.client.common_lib import error, barrier, logging_manager
 
 def deprecated(func):

@@ -278,6 +282,36 @@ def urlretrieve(url, filename, data=None, timeout=300):
         src_file.close()
 
 
+def hash(type, input=None):
+    """
+    Returns an hash object of type md5 or sha1. This function is implemented in
+    order to encapsulate hash objects in a way that is compatible with python
+    2.4 and python 2.6 without warnings.
+
+    Note that even though python 2.6 hashlib supports hash types other than
+    md5 and sha1, we are artificially limiting the input values in order to
+    make the function to behave exactly the same among both python
+    implementations.
+
+    @param input: Optional input string that will be used to update the hash.
+    """
+    if type not in ['md5', 'sha1']:
+        raise ValueError("Unsupported hash type: %s" % type)
+
+    try:
+        hash = hashlib.new(type)
+    except NameError:
+        if type == 'md5':
+            hash = md5.new()
+        elif type == 'sha1':
+            hash = sha.new()
+
+    if input:
+        hash.update(input)
+
+    return hash
+
+
 def get_file(src, dest, permissions=None):
     """Get a file from src, which can be local or a remote URL"""
     if src == dest:

index f8089f7..da7c543 100644
@@ -4,7 +4,7 @@ KVM test utility functions.
 @copyright: 2008-2009 Red Hat Inc.
 """
 
-import md5, sha, thread, subprocess, time, string, random, socket, os, signal
+import thread, subprocess, time, string, random, socket, os, signal
 import select, re, logging, commands, cPickle, pty
 from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import error

@@ -788,13 +788,7 @@ def hash_file(filename, size=None, method="md5"):
         size = fsize
     f = open(filename, 'rb')
 
-    if method == "md5":
-        hash = md5.new()
-    elif method == "sha1":
-        hash = sha.new()
-    else:
-        logging.error("Unknown hash type %s, returning None" % method)
-        return None
+    hash = utils.hash(method)
 
     while size > 0:
         if chunksize > size:

@@ -851,11 +845,7 @@ def unmap_url_cache(cachedir, url, expected_hash, method="md5"):
     failure_counter = 0
     while not file_hash == expected_hash:
         if os.path.isfile(file_local_path):
-            if method == "md5":
-                file_hash = hash_file(file_local_path, method="md5")
-            elif method == "sha1":
-                file_hash = hash_file(file_local_path, method="sha1")
-
+            file_hash = hash_file(file_local_path, method)
             if file_hash == expected_hash:
                 # File is already at the expected position and ready to go
                 src = file_from_url

index 8ff31da..90ff46d 100644
@@ -4,8 +4,8 @@ Utility functions to deal with ppm (qemu screendump format) files.
 @copyright: Red Hat 2008-2009
 """
 
-import md5, os, struct, time, re
-
+import os, struct, time, re
+from autotest_lib.client.bin import utils
 
 # Some directory/filename utils, for consistency
 

@@ -120,9 +120,9 @@ def image_md5sum(width, height, data):
     @data: PPM file data
     """
     header = "P6\n%d %d\n255\n" % (width, height)
-    md5obj = md5.new(header)
-    md5obj.update(data)
-    return md5obj.hexdigest()
+    hash = utils.hash('md5', header)
+    hash.update(data)
+    return hash.hexdigest()
 
 
 def get_region_md5sum(width, height, data, x1, y1, dx, dy,

index 456fb44..8ebe7c1 100644
@@ -4,7 +4,7 @@ Utilities to perform automatic guest installation using step files.
 @copyright: Red Hat 2008-2009
 """
 
-import os, time, md5, re, shutil, logging
+import os, time, re, shutil, logging
 from autotest_lib.client.common_lib import utils, error
 import kvm_utils, ppm_utils, kvm_subprocess
 try:

index 8b993a8..91b796d 100644
@@ -624,18 +624,3 @@ def interleave_entries(queue_entries, special_tasks):
 
     return interleaved_entries
 
-
-def get_sha1_hash(source):
-    """Gets the SHA-1 hash of the source string
-
-    @param source The string to hash
-    """
-    if sys.version_info < (2,5):
-        import sha
-        digest = sha.new()
-    else:
-        import hashlib
-        digest = hashlib.sha1()
-
-    digest.update(source)
-    return digest.hexdigest()

index c4d8988..7f62471 100644
@@ -3,7 +3,7 @@ import common
 from autotest_lib.frontend.afe import models as afe_models
 from autotest_lib.frontend.afe import model_logic, rpc_utils
 from autotest_lib.frontend.tko import models as tko_models
-from autotest_lib.client.common_lib import enum
+from autotest_lib.client.common_lib import enum, utils
 
 
 class Plan(dbmodels.Model):

@@ -102,7 +102,7 @@ class ControlFile(model_logic.ModelWithHash):
 
     @classmethod
     def _compute_hash(cls, **kwargs):
-        return rpc_utils.get_sha1_hash(kwargs['contents'])
+        return utils.hash('sha1', kwargs['contents']).hexdigest()
 
 
     def __unicode__(self):

@@ -322,8 +322,8 @@ class KeyVal(model_logic.ModelWithHash):
 
     @classmethod
     def _compute_hash(cls, **kwargs):
-        round1 = rpc_utils.get_sha1_hash(kwargs['key'])
-        return rpc_utils.get_sha1_hash(round1 + kwargs['value'])
+        round1 = utils.hash('sha1', kwargs['key']).hexdigest()
+        return utils.hash('sha1', round1 + kwargs['value']).hexdigest()
 
 
     def __unicode__(self):

index bc70074..8f06049 100644
@@ -1,4 +1,4 @@
-import os, md5
+import os
 
 from autotest_lib.client.common_lib import utils
 from autotest_lib.tko import utils as tko_utils

@@ -63,7 +63,7 @@ class kernel(object):
     @staticmethod
     def compute_hash(base, hashes):
         key_string = ','.join([base] + hashes)
-        return md5.new(key_string).hexdigest()
+        return utils.hash('md5', key_string).hexdigest()
 
 
 class test(object):

index 5110fe8..f0ccf55 100755
@@ -1,8 +1,9 @@
 #!/usr/bin/python
 
-import unittest, datetime, time, md5
+import unittest, datetime, time
 
 import common
+from autotest_lib.client.common_lib import utils
 from autotest_lib.tko.parsers import version_1
 
 

@@ -163,7 +164,7 @@ class test_status_line(unittest.TestCase):
                                       "patch0": "first_patch 0 0",
                                       "patch1": "another_patch 0 0"})
         kern = line.get_kernel()
-        kernel_hash = md5.new("2.6.24-rc40,0,0").hexdigest()
+        kernel_hash = utils.hash("2.6.24-rc40,0,0").hexdigest()
         self.assertEquals(kern.base, "2.6.24-rc40")
         self.assertEquals(kern.patches[0].spec, "first_patch")
         self.assertEquals(kern.patches[1].spec, "another_patch")

@@ -178,7 +179,7 @@ class test_status_line(unittest.TestCase):
                                       "patch0": "first_patch 0 0",
                                       "patch2": "another_patch 0 0"})
         kern = line.get_kernel()
-        kernel_hash = md5.new("2.6.24-rc40,0").hexdigest()
+        kernel_hash = utils.hash("2.6.24-rc40,0").hexdigest()
         self.assertEquals(kern.base, "2.6.24-rc40")
         self.assertEquals(kern.patches[0].spec, "first_patch")
         self.assertEquals(len(kern.patches), 1)

index d58975e..a5a797d 100755
@@ -12,10 +12,11 @@ Usage?  Just run it.
     utils/build_externals.py
 """
 
-import compileall, logging, os, sha, shutil, sys, tempfile, time, urllib2
+import compileall, logging, os, shutil, sys, tempfile, time, urllib2
 import subprocess, re
 import common
 from autotest_lib.client.common_lib import logging_config, logging_manager
+from autotest_lib.client.common_lib import utils
 
 # Where package source be fetched to relative to the top of the autotest tree.
 PACKAGE_DIR = 'ExternalSource'

@@ -152,7 +153,7 @@ def _checksum_file(full_path):
     """@returns The hex checksum of a file given its pathname."""
     inputfile = open(full_path, 'rb')
     try:
-        hex_sum = sha.sha(inputfile.read()).hexdigest()
+        hex_sum = utils.hash('sha1', inputfile.read()).hexdigest()
     finally:
         inputfile.close()
     return hex_sum

@@ -532,7 +533,7 @@ class ExternalPackage(object):
                 raise FetchError('%s from %s fails Content-Length %d '
                                  'sanity check.' % (self.name, url,
                                                     data_length))
-            checksum = sha.sha()
+            checksum = utils.hash('sha1')
             total_read = 0
             output = open(local_path, 'wb')
             try:

index b4f5646..4c1830e 100644
@@ -94,41 +94,92 @@ def extract_tarball(tarball):
         raise NameError('extracting tarball produced no dir')
 
 
-def get_md5sum(file_path):
-    """Gets the md5sum of a file. You must provide a valid path to the file"""
-    if not os.path.isfile(file_path):
-        raise ValueError, 'invalid file %s to verify' % file_path
-    md5sum = utils.system_output("md5sum " + file_path)
-    return md5sum.split()[0]
+def hash_file(filename, size=None, method="md5"):
+    """
+    Calculate the hash of filename.
+    If size is not None, limit to first size bytes.
+    Throw exception if something is wrong with filename.
+    Can be also implemented with bash one-liner (assuming size%1024==0):
+    dd if=filename bs=1024 count=size/1024 | sha1sum -
+
+    @param filename: Path of the file that will have its hash calculated.
+    @param method: Method used to calculate the hash. Supported methods:
+            * md5
+            * sha1
+    @returns: Hash of the file, if something goes wrong, return None.
+    """
+    chunksize = 4096
+    fsize = os.path.getsize(filename)
+
+    if not size or size > fsize:
+        size = fsize
+    f = open(filename, 'rb')
+
+    hash = utils.hash(method)
+
+    while size > 0:
+        if chunksize > size:
+            chunksize = size
+        data = f.read(chunksize)
+        if len(data) == 0:
+            logging.debug("Nothing left to read but size=%d" % size)
+            break
+        hash.update(data)
+        size -= len(data)
+    f.close()
+    return hash.hexdigest()
 
 
-def unmap_url_cache(cachedir, url, expected_md5):
+def unmap_url_cache(cachedir, url, expected_hash, method="md5"):
     """
     Downloads a file from a URL to a cache directory. If the file is already
-    at the expected position and has the expected md5 number, let's not
-    download it again.
+    at the expected position and has the expected hash, let's not download it
+    again.
+
+    @param cachedir: Directory that might hold a copy of the file we want to
+            download.
+    @param url: URL for the file we want to download.
+    @param expected_hash: Hash string that we expect the file downloaded to
+            have.
+    @param method: Method used to calculate the hash string (md5, sha1).
     """
     # Let's convert cachedir to a canonical path, if it's not already
     cachedir = os.path.realpath(cachedir)
     if not os.path.isdir(cachedir):
         try:
-            utils.system('mkdir -p ' + cachedir)
+            os.makedirs(cachedir)
         except:
             raise ValueError('Could not create cache directory %s' % cachedir)
     file_from_url = os.path.basename(url)
     file_local_path = os.path.join(cachedir, file_from_url)
-    if os.path.isfile(file_local_path):
-        file_md5 = get_md5sum(file_local_path)
-        if file_md5 == expected_md5:
-            # File is already at the expected position and ready to go
-            src = file_from_url
+
+    file_hash = None
+    failure_counter = 0
+    while not file_hash == expected_hash:
+        if os.path.isfile(file_local_path):
+            file_hash = hash_file(file_local_path, method)
+            if file_hash == expected_hash:
+                # File is already at the expected position and ready to go
+                src = file_from_url
+            else:
+                # Let's download the package again, it's corrupted...
+                logging.error("Seems that file %s is corrupted, trying to "
+                              "download it again" % file_from_url)
+                src = url
+                failure_counter += 1
         else:
-            # Let's download the package again, it's corrupted...
+            # File is not there, let's download it
             src = url
-    else:
-        # File is not there, let's download it
-        src = url
-    return utils.unmap_url(cachedir, src, cachedir)
+        if failure_counter > 1:
+            raise EnvironmentError("Consistently failed to download the "
+                                   "package %s. Aborting further download "
+                                   "attempts. This might mean either the "
+                                   "network connection has problems or the "
+                                   "expected hash string that was determined "
+                                   "for this file is wrong" % file_from_url)
+        file_path = utils.unmap_url(cachedir, src, cachedir)
+
+    return file_path
 
 
 def force_copy(src, dest):

index 483d71c..b8a03bf 100755
@@ -6,8 +6,9 @@ Program that calculates several hashes for a given CD image.
 """
 
 import os, sys, optparse, logging
-import common, kvm_utils
+import common
 from autotest_lib.client.common_lib import logging_config, logging_manager
+from autotest_lib.client.bin import utils
 
 
 class KvmLoggingConfig(logging_config.LoggingConfig):

@@ -44,11 +45,9 @@ if __name__ == "__main__":
         sys.exit(1)
 
     logging.info("Hash values for file %s", os.path.basename(filename))
-    logging.info("md5    (1m): %s", kvm_utils.hash_file(filename, 1024*1024,
-                                                        method="md5"))
-    logging.info("sha1   (1m): %s", kvm_utils.hash_file(filename, 1024*1024,
-                                                        method="sha1"))
-    logging.info("md5  (full): %s", kvm_utils.hash_file(filename,
-                                                        method="md5"))
-    logging.info("sha1 (full): %s", kvm_utils.hash_file(filename,
-                                                        method="sha1"))
+    logging.info("md5    (1m): %s", utils.hash_file(filename, 1024*1024,
+                                                    method="md5"))
+    logging.info("sha1   (1m): %s", utils.hash_file(filename, 1024*1024,
+                                                    method="sha1"))
+    logging.info("md5  (full): %s", utils.hash_file(filename, method="md5"))
+    logging.info("sha1 (full): %s", utils.hash_file(filename, method="sha1"))

100 lines not on this page · raw message

index 6731927..eb61ba2 100755
@@ -7,6 +7,7 @@ Utility classes and functions to handle Virtual Machine creation using qemu.
 
 import time, socket, os, logging, fcntl, re, commands
 import kvm_utils, kvm_subprocess
+from autotest_lib.client.bin import utils
 
 
 def get_image_filename(params, root_dir):

@@ -362,19 +363,19 @@ class VM:
             if params.get("md5sum_1m"):
                 logging.debug("Comparing expected MD5 sum with MD5 sum of "
                               "first MB of ISO file...")
-                actual_hash = kvm_utils.hash_file(iso, 1048576, method="md5")
+                actual_hash = utils.hash_file(iso, 1048576, method="md5")
                 expected_hash = params.get("md5sum_1m")
                 compare = True
             elif params.get("md5sum"):
                 logging.debug("Comparing expected MD5 sum with MD5 sum of ISO "
                               "file...")
-                actual_hash = kvm_utils.hash_file(iso, method="md5")
+                actual_hash = utils.hash_file(iso, method="md5")
                 expected_hash = params.get("md5sum")
                 compare = True
             elif params.get("sha1sum"):
                 logging.debug("Comparing expected SHA1 sum with SHA1 sum of "
                               "ISO file...")
-                actual_hash = kvm_utils.hash_file(iso, method="sha1")
+                actual_hash = utils.hash_file(iso, method="sha1")
                 expected_hash = params.get("sha1sum")
                 compare = True
             if compare: