diff mbox

[3/3] Client utils: Update utils.unmap_url_cache, added hash_file

Message ID 1265079896-5337-3-git-send-email-lmr@redhat.com (mailing list archive)
State New, archived
Headers show

Commit Message

Lucas Meneghel Rodrigues Feb. 2, 2010, 3:04 a.m. UTC
None
diff mbox

Patch

diff --git a/client/bin/base_utils.py b/client/bin/base_utils.py
index b4f5646..4c1830e 100644
--- a/client/bin/base_utils.py
+++ b/client/bin/base_utils.py
@@ -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):
diff --git a/client/tests/kvm/cd_hash.py b/client/tests/kvm/cd_hash.py
index 483d71c..b8a03bf 100755
--- a/client/tests/kvm/cd_hash.py
+++ b/client/tests/kvm/cd_hash.py
@@ -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"))
diff --git a/client/tests/kvm/kvm_utils.py b/client/tests/kvm/kvm_utils.py
index da7c543..cc2f678 100644
--- a/client/tests/kvm/kvm_utils.py
+++ b/client/tests/kvm/kvm_utils.py
@@ -767,42 +767,6 @@  def wait_for(func, timeout, first=0.0, step=1.0, text=None):
     return None
 
 
-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 get_hash_from_file(hash_path, dvd_basename):
     """
     Get the a hash from a given DVD image from a hash file
@@ -818,58 +782,6 @@  def get_hash_from_file(hash_path, dvd_basename):
             return line.split()[0]
 
 
-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 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:
-            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)
-
-    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:
-            # File is not there, let's download it
-            src = url
-        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 run_tests(test_list, job):
     """
     Runs the sequence of KVM tests based on the list of dictionaries
diff --git a/client/tests/kvm/kvm_vm.py b/client/tests/kvm/kvm_vm.py
index 6731927..eb61ba2 100755
--- a/client/tests/kvm/kvm_vm.py
+++ b/client/tests/kvm/kvm_vm.py
@@ -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: