X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/b13959674094c57ea0b5dce1164da4819bbcfe3c..9e4b4de20ef2c2c23a96fb1ac5d474fa2370695e:/image_creator/util.py diff --git a/image_creator/util.py b/image_creator/util.py index 3cacba0..61cc2a6 100644 --- a/image_creator/util.py +++ b/image_creator/util.py @@ -32,93 +32,52 @@ # or implied, of GRNET S.A. import sys -import pbs +import sh import hashlib -from clint.textui import colored -from progress.bar import Bar 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='', bar_type="default"): - - MESSAGE_LENGTH = 30 - - suffix={'default':'%(index)d/%(max)d', - 'percent':'%(percent)d%%', - 'b':'%(index)d/%(max)d B', - 'kb':'%(index)d/%(max)d KB', - 'mb':'%(index)d/%(max)d MB'} +class MD5: + def __init__(self, output): + self.out = output - return Bar(message=message.ljust(MESSAGE_LENGTH), fill='#', \ - suffix=suffix[bar_type]) + def compute(self, filename, size): -def md5(filename, size): + 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) - progressbar = progress("Calculating md5sum:", 'mb') - progressbar.max = (size // (2 ** 20)) - 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) // (2 ** 20)) - - checksum = md5.hexdigest() - output("\rCalculating md5sum...\033[K", False) - success(checksum) + checksum = md5.hexdigest() + progressbar.success(checksum) - return checksum + return checksum # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :