Revision e77e66a9 image_creator/util.py
| b/image_creator/util.py | ||
|---|---|---|
| 34 | 34 |
import sys |
| 35 | 35 |
import pbs |
| 36 | 36 |
import hashlib |
| 37 |
from colors import red, green, yellow |
|
| 38 |
from progress.bar import Bar |
|
| 39 | 37 |
|
| 40 | 38 |
|
| 41 | 39 |
class FatalError(Exception): |
| 42 | 40 |
pass |
| 43 | 41 |
|
| 44 | 42 |
|
| 45 |
silent = False |
|
| 46 |
|
|
| 47 |
|
|
| 48 | 43 |
def get_command(command): |
| 49 | 44 |
def find_sbin_command(command, exception): |
| 50 | 45 |
search_paths = ['/usr/local/sbin', '/usr/sbin', '/sbin'] |
| ... | ... | |
| 59 | 54 |
return find_sbin_command(command, e) |
| 60 | 55 |
|
| 61 | 56 |
|
| 62 |
def error(msg, new_line=True): |
|
| 63 |
nl = "\n" if new_line else '' |
|
| 64 |
sys.stderr.write(red('Error: %s' % msg) + nl)
|
|
| 65 |
|
|
| 66 |
|
|
| 67 |
def warn(msg, new_line=True): |
|
| 68 |
if not silent: |
|
| 69 |
nl = "\n" if new_line else '' |
|
| 70 |
sys.stderr.write(yellow("Warning: %s" % msg) + nl)
|
|
| 71 |
|
|
| 72 |
|
|
| 73 |
def success(msg, new_line=True): |
|
| 74 |
if not silent: |
|
| 75 |
nl = "\n" if new_line else '' |
|
| 76 |
sys.stdout.write(green(msg) + nl) |
|
| 77 |
if not nl: |
|
| 78 |
sys.stdout.flush() |
|
| 79 |
|
|
| 80 |
|
|
| 81 |
def output(msg="", new_line=True): |
|
| 82 |
if not silent: |
|
| 83 |
nl = "\n" if new_line else '' |
|
| 84 |
sys.stdout.write(msg + nl) |
|
| 85 |
if not nl: |
|
| 86 |
sys.stdout.flush() |
|
| 87 |
|
|
| 88 |
|
|
| 89 |
def progress(message='', bar_type="default"): |
|
| 90 |
|
|
| 91 |
MESSAGE_LENGTH = 30 |
|
| 92 |
|
|
| 93 |
suffix = {
|
|
| 94 |
'default': '%(index)d/%(max)d', |
|
| 95 |
'percent': '%(percent)d%%', |
|
| 96 |
'b': '%(index)d/%(max)d B', |
|
| 97 |
'kb': '%(index)d/%(max)d KB', |
|
| 98 |
'mb': '%(index)d/%(max)d MB' |
|
| 99 |
} |
|
| 100 |
|
|
| 101 |
bar = Bar() |
|
| 102 |
bar.message = message.ljust(MESSAGE_LENGTH) |
|
| 103 |
bar.fill = '#' |
|
| 104 |
bar.suffix = suffix[bar_type] |
|
| 105 |
bar.bar_prefix = ' [' |
|
| 106 |
bar.bar_suffix = '] ' |
|
| 107 |
|
|
| 108 |
return bar |
|
| 109 |
|
|
| 57 |
class MD5: |
|
| 58 |
def __init__(self, output): |
|
| 59 |
self.out = output |
|
| 110 | 60 |
|
| 111 |
def md5(filename, size):
|
|
| 61 |
def compute(self, filename, size):
|
|
| 112 | 62 |
|
| 113 |
BLOCKSIZE = 2 ** 22 # 4MB |
|
| 63 |
BLOCKSIZE = 2 ** 22 # 4MB
|
|
| 114 | 64 |
|
| 115 |
progressbar = progress("Calculating md5sum:", 'mb')
|
|
| 116 |
progressbar.max = ((size + 2 ** 20 - 1) // (2 ** 20)) |
|
| 117 |
md5 = hashlib.md5() |
|
| 118 |
with open(filename, "r") as src: |
|
| 119 |
left = size |
|
| 120 |
while left > 0: |
|
| 121 |
length = min(left, BLOCKSIZE) |
|
| 122 |
data = src.read(length) |
|
| 123 |
md5.update(data) |
|
| 124 |
left -= length |
|
| 125 |
progressbar.goto((size - left) // (2 ** 20)) |
|
| 65 |
progressbar = self.out.Progress("Calculating md5sum:", 'mb')
|
|
| 66 |
progressbar.max = ((size + 2 ** 20 - 1) // (2 ** 20))
|
|
| 67 |
md5 = hashlib.md5()
|
|
| 68 |
with open(filename, "r") as src:
|
|
| 69 |
left = size
|
|
| 70 |
while left > 0:
|
|
| 71 |
length = min(left, BLOCKSIZE)
|
|
| 72 |
data = src.read(length)
|
|
| 73 |
md5.update(data)
|
|
| 74 |
left -= length
|
|
| 75 |
progressbar.goto((size - left) // (2 ** 20))
|
|
| 126 | 76 |
|
| 127 |
checksum = md5.hexdigest() |
|
| 128 |
output("\rCalculating md5sum...\033[K", False)
|
|
| 129 |
success(checksum) |
|
| 77 |
checksum = md5.hexdigest() |
|
| 78 |
progressbar.success(checksum) |
|
| 130 | 79 |
|
| 131 |
return checksum |
|
| 80 |
return checksum
|
|
| 132 | 81 |
|
| 133 | 82 |
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai : |
Also available in: Unified diff