path: root/tools/mic

#!/usr/bin/env python
#
# Copyright (c) 2011 Intel, Inc.
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; version 2 of the License
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc., 59
# Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os, sys, errno
from mic import msger, creator
from mic.utils import cmdln, misc, errors
from mic.conf import configmgr
from mic.plugin import pluginmgr
from mic.__version__ import VERSION


def optparser_setup(func):
    """Setup optparser for a function"""
    if not hasattr(func, "optparser"):
        func.optparser = cmdln.SubCmdOptionParser()
        func.optparser.disable_interspersed_args()
    return func


class MicCmd(cmdln.Cmdln):
    """
    Usage: mic SUBCOMMAND [OPTS] [ARGS...]

    mic Means the Image Creation tool
    Try 'mic help SUBCOMMAND' for help on a specific subcommand.

    ${command_list}
    global ${option_list}
    ${help_list}
    """

    name = 'mic'
    version = VERSION

    def print_version(self):
        msger.raw("%s %s (%s)" % (self.name,
                                  self.version,
                                  misc.get_distro_str()))

    def get_optparser(self):
        optparser = cmdln.CmdlnOptionParser(self, version=self.version)
        # hook optparse print_version here
        optparser.print_version = self.print_version
        optparser.add_option('-d', '--debug', action='store_true',
                             dest='debug',
                             help='print debug message')
        optparser.add_option('-v', '--verbose', action='store_true',
                             dest='verbose',
                             help='verbose information')
        return optparser

    def postoptparse(self):
        if self.options.verbose:
            msger.set_loglevel('verbose')

        if self.options.debug:
            try:
                import rpm
                rpm.setVerbosity(rpm.RPMLOG_NOTICE)
            except ImportError:
                pass

            msger.set_loglevel('debug')

        self.print_version()

    def help_create(self):
        cr = creator.Creator()
        cr.optparser = cr.get_optparser()
        doc = cr.__doc__
        doc = cr._help_reindent(doc)
        doc = cr._help_preprocess(doc, None)
        doc = doc.replace(cr.name, "${cmd_name}", 1)
        doc = doc.rstrip() + '\n'
        return doc

    @cmdln.alias("cr")
    def do_create(self, argv):
        cr = creator.Creator()
        cr.main(argv[1:])

    def _root_confirm(self):
        if os.geteuid() != 0:
            msger.error('Root permission is required to continue, abort')

    @cmdln.alias("cv")
    @cmdln.option("-S", "--shell",
                  action="store_true", dest="shell", default=False,
                  help="Launch shell before packaging the converted image")
    def do_convert(self, subcmd, opts, *args):
        """${cmd_name}: convert image format

        Usage:
            mic convert <imagefile> <destformat>

        ${cmd_option_list}
        """

        if not args:
            # print help
            handler = self._get_cmd_handler('convert')
            if hasattr(handler, "optparser"):
                handler.optparser.print_help()
            return 1

        if len(args) == 1:
            raise errors.Usage("It need 2 arguments (1 given)")
        elif len(args) == 2:
            (srcimg, destformat) = args
        else:
            raise errors.Usage("Extra argument given")

        if not os.path.exists(srcimg):
            raise errors.CreatorError("Cannot find the image: %s" % srcimg)

        self._root_confirm()

        configmgr.convert['shell'] = opts.shell

        srcformat = misc.get_image_type(srcimg)
        if srcformat == "ext3fsimg":
            srcformat = "loop"

        srcimager = None
        destimager = None
        for iname, icls in pluginmgr.get_plugins('imager').iteritems():
            if iname == srcformat and hasattr(icls, "do_unpack"):
                srcimager = icls
            if iname == destformat and hasattr(icls, "do_pack"):
                destimager = icls

        if (srcimager and destimager) is None:
            raise errors.CreatorError("Can't convert from %s to %s" \
                                      % (srcformat, destformat))

        else:
            maptab = {
                        "livecd": "iso",
                        "liveusb": "usbimg",
                        "loop": "img",
                     }

            if destformat in maptab:
                imgname = os.path.splitext(os.path.basename(srcimg))[0]
                dstname = "{0}.{1}".format(imgname, maptab[destformat])

                if os.path.exists(dstname):
                    if msger.ask("Converted image %s seems existed, "
                                 "remove and continue?" % dstname):
                        os.unlink(dstname)
                    else:
                        raise errors.Abort("Canceled")

            base_on = srcimager.do_unpack(srcimg)
            destimager.do_pack(base_on)

    @cmdln.alias("ch")
    @cmdln.option('-s', '--saveto',
                  action='store', dest='saveto', default=None,
                  help="Save the unpacked image to specified dir")
    @optparser_setup
    def do_chroot(self, subcmd, opts, *args):
        """${cmd_name}: chroot into an image

        Usage:
            mic chroot [options] <imagefile> [command [arg]...]

        ${cmd_option_list}
        """

        if not args:
            # print help
            handler = self._get_cmd_handler('chroot')
            if hasattr(handler, "optparser"):
                handler.optparser.print_help()
            return 1

        targetimage = args[0]

        if not os.path.exists(targetimage):
            raise errors.CreatorError("Cannot find the image: %s"
                                      % targetimage)

        self._root_confirm()

        configmgr.chroot['saveto'] = opts.saveto

        imagetype = misc.get_image_type(targetimage)
        if imagetype in ("ext3fsimg", "ext4fsimg", "btrfsimg"):
            imagetype = "loop"

        chrootclass = None
        for pname, pcls in pluginmgr.get_plugins('imager').iteritems():
            if pname == imagetype and hasattr(pcls, "do_chroot"):
                chrootclass = pcls
                break

        if not chrootclass:
            raise errors.CreatorError("Cannot support image type: %s" \
                                      % imagetype)

        chrootclass.do_chroot(targetimage, args[1:])

if __name__ == "__main__":
    try:
        mic = MicCmd()
        sys.exit(mic.main())

    except KeyboardInterrupt:
        msger.error('\n^C catched, program aborted.')

    # catch 'no space left' exception, etc
    except IOError, e:
        if e.errno == errno.ENOSPC:
            msger.error('\nNo space left on device')
        raise

    except errors.Usage, usage:
        msger.error(str(usage))

    except errors.Abort, msg:
        msger.info(str(msg))

    except errors.CreatorError, err:
        if msger.get_loglevel() == 'debug':
            import traceback
            msger.error(traceback.format_exc())
        else:
            msger.error('\n'+str(err))