From c408053fabcc04a53440b0a2d602b89d7d22f1b1 Mon Sep 17 00:00:00 2001 From: Nikos Skalkotos Date: Wed, 29 Feb 2012 19:47:36 +0200 Subject: [PATCH] Add input options and populate image metadata Hooray! This is the first release that actually works! --- image_creator/disk.py | 15 ++++++- image_creator/main.py | 87 ++++++++++++++++++++++++++++++++----- image_creator/os_type/__init__.py | 7 +-- 3 files changed, 95 insertions(+), 14 deletions(-) diff --git a/image_creator/disk.py b/image_creator/disk.py index dd0127b..aebddc2 100644 --- a/image_creator/disk.py +++ b/image_creator/disk.py @@ -59,7 +59,7 @@ class Disk(object): def get_device(self): """Returns a newly created DiskDevice instance. - + This instance is a snapshot of the original source media of the Disk instance. """ @@ -195,7 +195,20 @@ class DiskDevice(object): start = last_partition['part_start'] / sector_size end = start + (block_size * block_cnt) / sector_size - 1 + self.g.part_del(dev, last_partition['part_num']) + self.g.part_add(dev, 'p', start, end) + return (end + 1) * sector_size + def size(self): + """Returns the "payload" size of the device. + + The size returned by this method is the size of the space occupied by + the partitions (including the space before the first partition). + """ + dev = self.g.part_to_dev(self.root) + last = self.g.part_list(dev)[-1] + + return last['part_end'] # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/image_creator/main.py b/image_creator/main.py index 9ff037e..6308a57 100644 --- a/image_creator/main.py +++ b/image_creator/main.py @@ -32,35 +32,102 @@ # or implied, of GRNET S.A. from image_creator import get_os_class +from image_creator import __version__ as version from image_creator.disk import Disk import sys import os +import optparse +from pbs import dd + + +class FatalError(Exception): + pass + + +def check_writable_dir(option, opt_str, value, parser): + if not os.path.isdir(value): + raise OptionValueError("%s is not a valid directory name" % value) + setattr(parser.values, option.dest, value) + + +def parse_options(input_args): + usage = "Usage: %prog [options] " + parser = optparse.OptionParser(version=version, usage=usage) + + parser.add_option("-o", "--outdir", type="string", dest="outdir", + default=".", action="callback", callback=check_writable_dir, + help="Output files to DIR [default: working dir]", + metavar="DIR") + + parser.add_option("-f", "--force", dest="force", default=False, + action="store_true", help="Overwrite output files if they exist") + + parser.add_option("--no-shrink", dest="shrink", default=True, + help="Don't shrink any partition before extracting the image", + action="store_false") + + options, args = parser.parse_args(input_args) + + if len(args) != 2: + parser.error('input media or name are missing') + options.source = args[0] + options.name = args[1] + + if not os.path.exists(options.source): + parser.error('Input media is not accessible') + + return options def main(): - if len(sys.argv) != 3: - sys.exit("Usage: %s " % - os.path.basename(sys.argv[0])) - source = sys.argv[1] - dest = sys.argv[2] - disk = Disk(source) + options = parse_options(sys.argv[1:]) + + if os.geteuid() != 0: + raise FatalError("You must run %s as root" \ + % os.path.basename(sys.argv[0])) + + if not options.force: + for ext in ('diskdump', 'meta'): + filename = "%s/%s.%s" % (options.outdir, options.name, ext) + if os.path.exists(filename): + raise FatalError("Output file %s exists " + "(use --force to overwrite it)." % filename) + + disk = Disk(options.source) try: dev = disk.get_device() dev.mount() osclass = get_os_class(dev.distro, dev.ostype) image_os = osclass(dev.root, dev.g) metadata = image_os.get_metadata() - for key in metadata.keys(): - print "%s=%s" % (key, metadata[key]) image_os.data_cleanup() dev.umount() - #dev.shrink() + size = options.shrink and dev.shrink() or dev.size() + metadata['size'] = str(size // 2 ** 20) + dd('if=%s' % dev.device, + 'of=%s/%s.%s' % (options.outdir, options.name, 'diskdump'), + 'bs=4M', 'count=%d' % ((size + 1) // 2 ** 22)) + + f = open('%s/%s.%s' % (options.outdir, options.name, 'meta'), 'w') + for key in metadata.keys(): + f.write("%s=%s\n" % (key, metadata[key])) + f.close() finally: disk.cleanup() + return 0 + +COLOR_BLACK = "\033[00m" +COLOR_RED = "\033[1;31m" + if __name__ == '__main__': - main() + try: + ret = main() + sys.exit(ret) + except FatalError as e: + print >> sys.stderr, "\n%sError: %s%s\n" % (COLOR_RED, e, COLOR_BLACK) + sys.exit(1) # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py index f1f74ff..5fd00ab 100644 --- a/image_creator/os_type/__init__.py +++ b/image_creator/os_type/__init__.py @@ -71,9 +71,10 @@ class OSBase(object): def get_metadata(self): """Returns some descriptive metadata about the OS.""" meta = {} - meta["OSFAMILY"] = self.g.inspect_get_type(self.root) - meta["OS"] = self.g.inspect_get_distro(self.root) - meta["description"] = self.g.inspect_get_product_name(self.root) + meta['ROOT_PARTITION'] = "%d" % self.g.part_to_partnum(self.root) + meta['OSFAMILY'] = self.g.inspect_get_type(self.root) + meta['OS'] = self.g.inspect_get_distro(self.root) + meta['description'] = self.g.inspect_get_product_name(self.root) return meta -- 1.7.10.4