Synnefo » snf-image-creator

#!/usr/bin/env python

# -*- coding: utf-8 -*-

#

# Copyright 2012 GRNET S.A. All rights reserved.

#

# Redistribution and use in source and binary forms, with or

# without modification, are permitted provided that the following

# conditions are met:

#

#   1. Redistributions of source code must retain the above

#      copyright notice, this list of conditions and the following

#      disclaimer.

#

#   2. Redistributions in binary form must reproduce the above

#      copyright notice, this list of conditions and the following

#      disclaimer in the documentation and/or other materials

#      provided with the distribution.

#

# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS

# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED

# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR

# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR

# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,

# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT

# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF

# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED

# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT

# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN

# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE

# POSSIBILITY OF SUCH DAMAGE.

#

# The views and conclusions contained in the software and

# documentation are those of the authors and should not be

# interpreted as representing official policies, either expressed

# or implied, of GRNET S.A.

"""This module is the entrance point for the dialog-based version of the

snf-image-creator program. The main function will create a dialog where the

user is asked if he wants to use the program in expert or wizard mode.

"""

import dialog

import sys

import os

import stat

import textwrap

import signal

import optparse

import types

from image_creator import __version__ as version

from image_creator.util import FatalError

from image_creator.output import Output

from image_creator.output.cli import SimpleOutput

from image_creator.output.dialog import GaugeOutput

from image_creator.output.composite import CompositeOutput

from image_creator.disk import Disk

from image_creator.dialog_wizard import start_wizard

from image_creator.dialog_menu import main_menu

from image_creator.dialog_util import SMALL_WIDTH, WIDTH, confirm_exit, \

    Reset, update_background_title

def create_image(d, media, out, tmp):

    """Create an image out of `media'"""

    d.setBackgroundTitle('snf-image-creator')

    gauge = GaugeOutput(d, "Initialization", "Initializing...")

    out.add(gauge)

    disk = Disk(media, out, tmp)

    def signal_handler(signum, frame):

        gauge.cleanup()

        disk.cleanup()

    signal.signal(signal.SIGINT, signal_handler)

    signal.signal(signal.SIGTERM, signal_handler)

    try:

        snapshot = disk.snapshot()

        image = disk.get_image(snapshot)

        out.output("Collecting image metadata ...")

        metadata = {}

        for (key, value) in image.meta.items():

            metadata[str(key)] = str(value)

        for (key, value) in image.os.meta.items():

            metadata[str(key)] = str(value)

        out.success("done")

        gauge.cleanup()

        out.remove(gauge)

        # Make sure the signal handler does not call gauge.cleanup again

        def dummy(self):

            pass

        gauge.cleanup = type(GaugeOutput.cleanup)(dummy, gauge, GaugeOutput)

        session = {"dialog": d,

                   "disk": disk,

                   "image": image,

                   "metadata": metadata}

        msg = "snf-image-creator detected a %s system on the input media. " \

              "Would you like to run a wizard to assist you through the " \

              "image creation process?\n\nChoose <Wizard> to run the wizard," \

              " <Expert> to run the snf-image-creator in expert mode or " \

              "press ESC to quit the program." \

              % (image.ostype if image.ostype == image.distro or

                 image.distro == "unknown" else "%s (%s)" %

                 (image.ostype, image.distro))

        update_background_title(session)

        while True:

            code = d.yesno(msg, width=WIDTH, height=12, yes_label="Wizard",

                           no_label="Expert")

            if code == d.DIALOG_OK:

                if start_wizard(session):

                    break

            elif code == d.DIALOG_CANCEL:

                main_menu(session)

                break

            if confirm_exit(d):

                break

        d.infobox("Thank you for using snf-image-creator. Bye", width=53)

    finally:

        disk.cleanup()

    return 0

def select_file(d, media):

    """Select a media file"""

    if media == '/':

        return '/'

    default = os.getcwd() + os.sep

    while 1:

        if media is not None:

            if not os.path.exists(media):

                d.msgbox("The file `%s' you choose does not exist." % media,

                         width=SMALL_WIDTH)

            else:

                mode = os.stat(media).st_mode

                if not stat.S_ISDIR(mode):

                    break

                default = media

        (code, media) = d.fselect(default, 10, 60, extra_button=1,

                                  title="Please select an input media.",

                                  extra_label="Bundle Host")

        if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):

            if confirm_exit(d, "You canceled the media selection dialog box."):

                sys.exit(0)

            else:

                media = None

                continue

        elif code == d.DIALOG_EXTRA:

            return '/'

    return media

def _dialog_form(self, text, height=20, width=60, form_height=15, fields=[],

                 **kwargs):

    """Display a form box.

    fields is in the form: [(label1, item1, item_length1), ...]

    """

    cmd = ["--form", text, str(height), str(width), str(form_height)]

    label_len = 0

    for field in fields:

        if len(field[0]) > label_len:

            label_len = len(field[0])

    input_len = width - label_len - 1

    line = 1

    for field in fields:

        label = field[0]

        item = field[1]

        item_len = field[2]

        cmd.extend((label, str(line), str(1), item, str(line),

                   str(label_len + 1), str(input_len), str(item_len)))

        line += 1

    code, output = self._perform(*(cmd,), **kwargs)

    if not output:

        return (code, [])

    return (code, output.splitlines())

def main():

    d = dialog.Dialog(dialog="dialog")

    # Add extra button in dialog library

    dialog._common_args_syntax["extra_button"] = \

        lambda enable: dialog._simple_option("--extra-button", enable)

    dialog._common_args_syntax["extra_label"] = \

        lambda string: ("--extra-label", string)

    # Allow yes-no label overwriting

    dialog._common_args_syntax["yes_label"] = \

        lambda string: ("--yes-label", string)

    dialog._common_args_syntax["no_label"] = \

        lambda string: ("--no-label", string)

    # Monkey-patch pythondialog to include support for form dialog boxes

    if not hasattr(dialog, 'form'):

        d.form = types.MethodType(_dialog_form, d)

    usage = "Usage: %prog [options] [<input_media>]"

    parser = optparse.OptionParser(version=version, usage=usage)

    parser.add_option("-l", "--logfile", type="string", dest="logfile",

                      default=None, help="log all messages to FILE",

                      metavar="FILE")

    parser.add_option("--tmpdir", type="string", dest="tmp", default=None,

                      help="create large temporary image files under DIR",

                      metavar="DIR")

    options, args = parser.parse_args(sys.argv[1:])

    if len(args) > 1:

        parser.error("Wrong number of arguments")

    d.setBackgroundTitle('snf-image-creator')

    try:

        if os.geteuid() != 0:

            raise FatalError("You must run %s as root" %

                             parser.get_prog_name())

        if options.tmp is not None and not os.path.isdir(options.tmp):

            raise FatalError("The directory `%s' specified with --tmpdir is "

                             "not valid" % options.tmp)

        logfile = None

        if options.logfile is not None:

            try:

                logfile = open(options.logfile, 'w')

            except IOError as e:

                raise FatalError(

                    "Unable to open logfile `%s' for writing. Reason: %s" %

                    (options.logfile, e.strerror))

        media = select_file(d, args[0] if len(args) == 1 else None)

        try:

            log = SimpleOutput(False, logfile) if logfile is not None \

                else Output()

            while 1:

                try:

                    out = CompositeOutput([log])

                    out.output("Starting %s v%s ..." %

                               (parser.get_prog_name(), version))

                    ret = create_image(d, media, out, options.tmp)

                    sys.exit(ret)

                except Reset:

                    log.output("Resetting everything ...")

                    continue

        finally:

            if logfile is not None:

                logfile.close()

    except FatalError as e:

        msg = textwrap.fill(str(e), width=WIDTH)

        d.infobox(msg, width=WIDTH, title="Fatal Error")

        sys.exit(1)

# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :