From d603d80deb738690add1cf100118760ef77ce853 Mon Sep 17 00:00:00 2001 From: Nikos Skalkotos Date: Fri, 23 Mar 2012 19:16:23 +0200 Subject: [PATCH] Move image_extract method into DiskDevice --- image_creator/disk.py | 53 ++++++++++++++++++++++++++++++++++++++++--------- image_creator/main.py | 32 ++--------------------------- image_creator/util.py | 1 - 3 files changed, 46 insertions(+), 40 deletions(-) diff --git a/image_creator/disk.py b/image_creator/disk.py index 7fc0cdf..f78adc5 100644 --- a/image_creator/disk.py +++ b/image_creator/disk.py @@ -31,9 +31,9 @@ # interpreted as representing official policies, either expressed # or implied, of GRNET S.A. -from image_creator.util import get_command, warn, progress_generator +from image_creator.util import get_command, warn, progress_generator, success from image_creator import FatalError -from clint.textui import indent, puts, colored +from clint.textui import puts import stat import os @@ -43,6 +43,7 @@ import re import sys import guestfs import time +from sendfile import sendfile class DiskError(Exception): @@ -104,16 +105,16 @@ class Disk(object): sourcedev = self.source mode = os.stat(self.source).st_mode if stat.S_ISDIR(mode): - puts(colored.green('looks like a directory')) + success('looks like a directory') return self._losetup(self._dir_to_disk()) elif stat.S_ISREG(mode): - puts(colored.green('looks like an image file')) + success('looks like an image file') sourcedev = self._losetup(self.source) elif not stat.S_ISBLK(mode): raise ValueError("Invalid media source. Only block devices, " "regular files and directories are supported.") else: - puts(colored.green('looks like a block device')) + success('looks like a block device') # Take a snapshot and return it to the user puts("Snapshotting media source...", False) @@ -140,7 +141,7 @@ class Disk(object): finally: os.unlink(table) - puts(colored.green('done')) + success('done') new_device = DiskDevice("/dev/mapper/%s" % snapshot) self._devices.append(new_device) new_device.enable() @@ -197,7 +198,7 @@ class DiskDevice(object): self.root = roots[0] self.ostype = self.g.inspect_get_type(self.root) self.distro = self.g.inspect_get_distro(self.root) - puts(colored.green('found a %s system' % self.distro)) + success('found a %s system' % self.distro) def destroy(self): """Destroy this DiskDevice instance.""" @@ -250,6 +251,7 @@ class DiskDevice(object): (in bytes) is returned. """ puts("Shrinking image (this may take a while)...", False) + sys.stdout.flush() dev = self.g.part_to_dev(self.root) parttype = self.g.part_get_parttype(dev) @@ -287,8 +289,8 @@ class DiskDevice(object): self.g.part_add(dev, 'p', start, end) new_size = (end + 1) * sector_size - puts(colored.green("new image size is %dMB\n" % (new_size // 2 ** 20))) - + success("new image size is %dMB" % + ((new_size + 2 ** 20 - 1) // 2 ** 20)) return new_size def size(self): @@ -302,4 +304,37 @@ class DiskDevice(object): return last['part_end'] + 1 + def dump(self, outfile): + """Dumps the content of device into a file. + + This method will only dump the actual payload, found by reading the + partition table. Empty space in the end of the device will be ignored. + """ + blocksize = 2 ** 22 # 4MB + size = self.size() + progress_size = (size + 2 ** 20 - 1) // 2 ** 20 # in MB + progressbar = progress_generator("Dumping image file: ", progress_size) + + source = open(self.device, "r") + try: + dest = open(outfile, "w") + try: + left = size + offset = 0 + progressbar.next() + while left > 0: + length = min(left, blocksize) + sent = sendfile(dest.fileno(), source.fileno(), offset, + length) + offset += sent + left -= sent + for i in range((length + 2 ** 20 - 1) // 2 ** 20): + progressbar.next() + finally: + dest.close() + finally: + source.close() + + success('Image file %s was successfully created' % outfile) + # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/image_creator/main.py b/image_creator/main.py index 40bc880..7a2a4fe 100755 --- a/image_creator/main.py +++ b/image_creator/main.py @@ -38,8 +38,7 @@ from image_creator import __version__ as version from image_creator import FatalError from image_creator.disk import Disk from image_creator.util import get_command, error, progress_generator, success -from clint.textui import puts, indent -from sendfile import sendfile +from clint.textui import puts import sys import os @@ -108,33 +107,6 @@ def parse_options(input_args): return options -def extract_image(device, outfile, size): - blocksize = 4194304 # 4MB - progress_size = (size + 1048575) // 1048576 # in MB - progressbar = progress_generator("Dumping image file: ", - progress_size) - source = open(device, "r") - try: - dest = open(outfile, "w") - try: - left = size - offset = 0 - progressbar.next() - while left > 0: - length = min(left, blocksize) - sent = sendfile(dest.fileno(), source.fileno(), offset, length) - offset += sent - left -= sent - for i in range(4): - progressbar.next() - finally: - dest.close() - finally: - source.close() - - success('Image file %s was successfully created' % outfile) - - def image_creator(): puts('snf-image-creator %s\n' % version) options = parse_options(sys.argv[1:]) @@ -180,7 +152,7 @@ def image_creator(): finally: f.close() - extract_image(dev.device, options.outfile, size) + dev.dump(options.outfile) finally: puts('cleaning up...') disk.cleanup() diff --git a/image_creator/util.py b/image_creator/util.py index a0d44ba..50304f0 100644 --- a/image_creator/util.py +++ b/image_creator/util.py @@ -69,5 +69,4 @@ def progress_generator(label='', n=100): position = yield yield # suppress the StopIteration exception - # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : -- 1.7.10.4