#!/usr/bin/python -t
#
# mic-chroot : chroot into an image or file system in order to change it
#
# Copyright 2009, 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 Library 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
import os.path
import sys
import optparse
import shutil
import subprocess
import signal

import mic.imgcreate as imgcreate
import mic.imgconvert as imgconvert
import mic.chroot as chroot

def sig_interrupt(signum=None, frame=None):
    raise imgcreate.CreatorError('Canceled by ^C.')
signal.signal(signal.SIGINT, sig_interrupt)

try:
    import mic.__version__
    version = mic.__version__.version
except:
    version = 'unknown'

class Usage(Exception):
    def __init__(self, msg = None, no_error = False):
        Exception.__init__(self, msg, no_error)

def parse_options(args):
    parser = optparse.OptionParser(
        usage="mic-chroot [options] <image | filesystem>\n\n"
              "<image> is a live image file name, <filesystem>"
              " is a directory path which\n contains the whole file system"
              " unpacked from a live image."
        , version = version)

    parser.add_option("-s", "--save-to", type="string", dest="saveto",
                    help="Save unpacked filesystem to the specified path ")
    parser.add_option("", "--unpack-only", action="store_true",
                      dest="unpackonly", default=False,
                      help="Just unpack an image, this is used to"
                           "unpack an image with -s option together"
                     )
    parser.add_option("-b", "--bind-mounts", type="string", dest="bindmounts",
                      help="Specify bind mount list, for example: -b "
                           "\"/proc:/proc;/:/parentroot\""
                     )
    parser.add_option("-c", "--convert-to", type="string", dest="convertto",
                      help="Convert it to the specified type live image on"
                           " exiting, the allowed value is livecd or liveusb"
                     )
    parser.add_option("-e", "--execute", type="string", dest="execute",
                      help="Execute the given command within the chroot"
                           " instead of an interactive shell"
                     )
    parser.add_option("", "--convert-only", action="store_true",
                      dest="convertonly", default=False,
                      help="Just convert an image, this will skip chroot and"
                           " directly convert an image/filesytem with -c "
                           "option together"
                     )
    parser.add_option("-o", "--outdir", type="string",
                      dest="outdir", default=None,
                      help="Output directory to use "
                           "(default: current work dir)"
                     )
    imgcreate.setup_logging(parser)

    (options, args) = parser.parse_args()

    if len(args) != 1:
        parser.print_help()
        sys.exit(-1)

    target = args[0]
    if os.path.isfile(target):
        type = imgcreate.get_image_type(target)
        if type not in ("livecd", "liveusb", "ext3fsimg", "raw"):
            chroot.perror("I can't recognize this image type.")
            return None
        options.targettype = "img"
    elif os.path.isdir(target):
        options.targettype = "dir"
    else:
        chroot.perror("%s doesn't exist." % target)
        return None

    options.targetpath = target
    if options.convertto:
        if options.convertto not in ("livecd", "liveusb", "ext3fsimg"):
            chroot.perror("I can't convert it to %s" % options.convertto)
            return None

    if not options.unpackonly and options.bindmounts:
        if chroot.check_bind_mounts(None, options.bindmounts) == False:
            chroot.perror("bindmounts %s is invalid." % options.bindmounts)
            return None

    return options

def print_danger_warning():
    print "+====================================================================================+"
    print "| WARNING: this script is very dangerous, it may damage your system data,            |"
    print "| so please use it carefully and read this warning seriously.                        |"
    print "|                                                                                    |"
    print "| mic-chroot will create some temporary directories under /var/tmp, their names      |"
    print "| look like 'mic-tmp-*', mic-chroot also creates a directory to save image           |"
    print "| file system for chroot. For such directories, please use mic-rm-chroot-dir to      |"
    print "| remove, otherwise, you may remove your whole system by 'rm -rf <yourdir>' if you   |"
    print "| terminate mic-chroot because of some exceptions or unknown errors, the root        |"
    print "| cause is mic-chroot will bind mount / to chroot directory, exceptional             |"
    print "| terminations didn't unmount it, so you may damage your system, we have warned you. |"
    print "| Please use it very very carefully, and don't use 'rm -rf' for such directories,    |"
    print "| instead, use 'mic-rm-chroot-dir'                                                   |"
    print "+====================================================================================+"
    print "\n"

def cleanup_after_chroot(targettype,imgmount,tmpdir,tmpmnt):
    if imgmount and targettype == "img":
        imgmount.cleanup()
    if tmpdir:
        shutil.rmtree(tmpdir, ignore_errors = True)
    if tmpmnt:
        shutil.rmtree(tmpmnt, ignore_errors = True)    

def main():
    try:
        options = parse_options(sys.argv[1:])
        if not options:
            return 2
    except Usage, (msg, no_error):
        if no_error:
            out = sys.stdout
            ret = 0
        else:
            out = sys.stderr
            ret = 2
        if msg:
            print >> out, msg
        return ret

    if os.geteuid () != 0:
        chroot.perror("You must run mic-chroot as root")
        return 1

    print_danger_warning()

    imgmount = None
    tmpdir = None
    tmpmnt = None
    if options.outdir:
        if not os.path.isdir(options.outdir):
            chroot.perror("outdir %s doesn't exist." % options.outdir)
            return 1
    else:
       options.outdir = "."

    if options.saveto:
        if os.path.exists(options.saveto):
            chroot.perror("saveto path %s has existed." % options.saveto)
            return 1
        else:
            imgcreate.makedirs(options.saveto)
        savefs = True
    else:
        if options.unpackonly:
            chroot.perror("Please give out your saveto dir.")
            return 1
        if os.path.isdir(options.targetpath):
            options.saveto = options.targetpath
        else:
            options.saveto = imgcreate.mkdtemp()
            tmpdir = options.saveto
            chroot.pinfo("%s will be unpacked into %s" % (options.targetpath, options.saveto))
        savefs = False
    options.saveto = os.path.abspath(os.path.expanduser(options.saveto))

    if options.targettype == "img":
        imgmount = chroot.ImageMount(options.targetpath)
        try:
            if savefs:
                tmpmnt = imgcreate.mkdtemp()
                imgmount.mount(tmpmnt)
                print("Copying filesystem to %s..." % options.saveto)
                imgcreate.myxcopytree(tmpmnt, options.saveto)
            else:
                imgmount.mount(options.saveto)
        except imgcreate.CreatorError, e:
            cleanup_after_chroot(options.targettype,imgmount,tmpdir,tmpmnt)
            chroot.perror("ImageMount: %s" % e)
            return 1

    if options.unpackonly:
        imgmount.cleanup()
        if tmpmnt:
            shutil.rmtree(tmpmnt, ignore_errors = True)
        return 0

    if options.bindmounts:
        if chroot.check_bind_mounts(options.saveto, 
                                   options.bindmounts) == False:
            cleanup_after_chroot(options.targettype,imgmount,tmpdir,tmpmnt)
            chroot.perror("bindmounts %s is invalid." % options.bindmounts)
            return 1

    if not options.convertonly:
        if not options.execute:
            # If executable was not given we use following defaults
            options.execute = "/bin/env HOME=/root /bin/bash"

        try:
            chroot.chroot(options.saveto, bindmounts = options.bindmounts, execute = options.execute)
        finally:
            cleanup_after_chroot(options.targettype,imgmount,tmpdir,tmpmnt)
        
    if options.convertto:
        chroot.cleanup_mounts(options.saveto)
        if options.targettype == "img":
            args = ["mic-image-convertor", "--source-image=%s" % imgmount.os_image,
                    "--target-format=%s" % options.convertto, "--outdir=%s" % options.outdir]
        else:
            args = ["mic-image-convertor", "--source-image=%s" % options.saveto,
                    "--target-format=%s" % options.convertto, "--outdir=%s" % options.outdir]
        subprocess.call(args)
    if options.convertonly and not options.convertto:
        chroot.perror("If --convert-only is given, you must use -c | --convert-to to specify target image format.")
        return 1
    
    cleanup_after_chroot(options.targettype,imgmount,tmpdir,tmpmnt)
    print "Finished."
    return 0

COLOR_BLACK = "\033[00m"
COLOR_RED =   "\033[1;31m"
if __name__ == "__main__":
    try:
        sys.exit(main())
    except (imgcreate.CreatorError), e:
        print >> sys.stderr, "\n%sError: %s%s\n" % (COLOR_RED, e, COLOR_BLACK)
        imgcreate.terminate_log_thread()
        sys.exit(1)
        
