X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/8e3065a039b08d9c37d183f8cc3bb8ae5dcd4410..6523456eff0071b3e87af2b7ee7465ebf49d527b:/image_creator/util.py?ds=sidebyside diff --git a/image_creator/util.py b/image_creator/util.py index 2805de7..61cc2a6 100644 --- a/image_creator/util.py +++ b/image_creator/util.py @@ -32,87 +32,52 @@ # or implied, of GRNET S.A. import sys -import pbs +import sh import hashlib -from clint.textui import colored, progress as uiprogress class FatalError(Exception): pass -silent = False - - def get_command(command): def find_sbin_command(command, exception): 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(colored.red('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() - - -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 progress(message=''): - - PROGRESS_LENGTH = 32 - MESSAGE_LENGTH = 32 +class MD5: + def __init__(self, output): + self.out = output - def progress_generator(n=100): - position = 0 - msg = message.ljust(MESSAGE_LENGTH) - for i in uiprogress.bar(range(n), msg, PROGRESS_LENGTH, silent): - if i < position: - continue - position = yield - yield # suppress the StopIteration exception - return progress_generator + def compute(self, filename, size): -def md5(filename, size, progress = None): + MB = 2 ** 20 + BLOCKSIZE = 4 * MB # 4MB - BLOCKSIZE = 2^22 # 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) - 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 + checksum = md5.hexdigest() + progressbar.success(checksum) - return md5.hexdigest() + return checksum # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :