X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/979096ddfddb0cc13151c7e223be64e80ec8db44..6523456eff0071b3e87af2b7ee7465ebf49d527b:/image_creator/util.py diff --git a/image_creator/util.py b/image_creator/util.py index 588d2b0..61cc2a6 100644 --- a/image_creator/util.py +++ b/image_creator/util.py @@ -32,10 +32,12 @@ # or implied, of GRNET S.A. import sys -import pbs -from clint.textui import colored, progress as uiprogress +import sh +import hashlib -silent = False + +class FatalError(Exception): + pass def get_command(command): @@ -43,55 +45,39 @@ def get_command(command): search_paths = ['/usr/local/sbin', '/usr/sbin', '/sbin'] for fullpath in map(lambda x: "%s/%s" % (x, command), search_paths): if os.path.exists(fullpath) and os.access(fullpath, os.X_OK): - return pbs.Command(fullpath) + return sh.Command(fullpath) raise exception try: - return pbs.__getattr__(command) - except pbs.CommadNotFount as e: + return sh.__getattr__(command) + except sh.CommadNotFount as e: return find_sbin_command(command, e) -def error(msg, new_line=True): - nl = "\n" if new_line else '' - sys.stderr.write('Error: %s' % msg + nl) - - -def warn(msg, new_line=True): - if not silent: - nl = "\n" if new_line else '' - sys.stderr.write(colored.yellow("Warning: %s" % msg) + nl) - - -def success(msg, new_line=True): - if not silent: - nl = "\n" if new_line else '' - sys.stdout.write(colored.green(msg) + nl) - if not nl: - sys.stdout.flush() - +class MD5: + def __init__(self, output): + self.out = output -def output(msg="", new_line=True): - if not silent: - nl = "\n" if new_line else '' - sys.stdout.write(msg + nl) - if not nl: - sys.stdout.flush() + def compute(self, filename, size): + MB = 2 ** 20 + BLOCKSIZE = 4 * MB # 4MB -def progress(label='', n=100): + 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) - PROGRESS_LENGTH = 32 - MESSAGE_LENGTH = 32 + checksum = md5.hexdigest() + progressbar.success(checksum) - def progress_generator(label, n): - position = 0 - for i in uiprogress.bar(range(n), label.ljust(MESSAGE_LENGTH), \ - PROGRESS_LENGTH, silent): - if i < position: - continue - position = yield - yield # suppress the StopIteration exception - return progress_generator(label, n) + return checksum # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :