atomic

#!/usr/bin/python -Es
# Copyright (C) 2014-2015 Red Hat
# AUTHOR: Dan Walsh <dwalsh@redhat.com>
# see file 'COPYING' for use and warranty information
#
# atomic is a tool for managing Atomic Systems and Containers
#
#    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; either version 2 of
#    the License, or (at your option) any later version.
#
#    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., 51 Franklin Street, Fifth Floor, Boston, MA
#    02110-1301 USA.
#
#
import os
import sys
import gettext
import argparse
import subprocess

import docker
import Atomic
from Atomic import diff
from Atomic import images
from Atomic import scan
from Atomic import top
from Atomic import mount
from Atomic import verify
from Atomic import storage
from Atomic import host
from Atomic import info
from Atomic import install
from Atomic import help as Help
from Atomic import run
from Atomic import containers
from Atomic import pull
from Atomic import push
from Atomic import stop
from Atomic import uninstall
from Atomic import update
from Atomic.util import write_err, NoDockerDaemon
import traceback
from Atomic.mount import MountError

PROGNAME = "atomic"
gettext.bindtextdomain(PROGNAME, "/usr/share/locale")
gettext.textdomain(PROGNAME)
try:
    # pylint: disable=unexpected-keyword-arg
    gettext.install(PROGNAME, unicode=True, codeset='utf-8')
except TypeError:
    # Failover to python3 install
    gettext.install(PROGNAME, codeset='utf-8')
except IOError:
    import builtins # pylint: disable=import-error
    builtins.__dict__['_'] = str

class HelpByDefaultArgumentParser(argparse.ArgumentParser):

    def error(self, message):
        write_err('%s: %s' % (self.prog, message))
        write_err("Try '%s --help' for more information." % self.prog)
        sys.exit(2)

    def print_usage(self, message="too few arguments"):
        self.prog = " ".join(sys.argv)
        self.error(message)

# Code for python2 copied from Adrian Sampson hack
# https://gist.github.com/sampsyo/471779
#
class AliasedSubParsersAction(argparse._SubParsersAction): # pylint: disable=protected-access
    class _AliasedPseudoAction(argparse.Action): # pylint: disable=abstract-method
        def __init__(self, name, aliases, help_text):
            dest = name
            if aliases:
                dest += ' (%s)' % ','.join(aliases)
            sup = super(AliasedSubParsersAction._AliasedPseudoAction, self) #pylint: disable=protected-access
            sup.__init__(option_strings=[], dest=dest, help=help_text)

    def add_parser(self, name, **kwargs):
        if 'aliases' in kwargs:
            aliases = kwargs['aliases']
            del kwargs['aliases']
        else:
            aliases = []

        parser = super(AliasedSubParsersAction, self).add_parser(name, **kwargs)

        # Make the aliases work.
        for alias in aliases:
            self._name_parser_map[alias] = parser
        # Make the help text reflect them, first removing old help entry.
        if 'help' in kwargs:
            help = kwargs.pop('help') # pylint: disable=redefined-builtin
            self._choices_actions.pop()
            pseudo_action = self._AliasedPseudoAction(name, aliases, help)
            self._choices_actions.append(pseudo_action)

        return parser

def need_root():
    sub_function = sys.argv[1] if sys.argv[1] not in ['--debug'] else sys.argv[2]
    exit("Some operations for '%s' require root access." % sub_function)
    sys.exit(1)

def create_parser(help_text):
    parser = HelpByDefaultArgumentParser(description=help_text)
    parser.register('action', 'parsers', AliasedSubParsersAction)
    parser.add_argument('-v', '--version', action='version', version=Atomic.__version__)
    parser.add_argument('--debug', default=False, action='store_true')
    subparser = parser.add_subparsers(help=_("commands"))

    # atomic cluster (if commctl.cli is available)
    try:
        import commctl.cli
        clusterp = subparser.add_parser(
            "cluster", help=_("execute cluster management commands"))
        commctl.cli.add_cluster_commands(clusterp)
    except ImportError:
        pass

    diff.cli(subparser)
    Help.cli(subparser)

    if os.path.exists("/usr/bin/rpm-ostree") and os.path.exists("/run/ostree-booted"):
        host.cli(subparser)
    info.cli(subparser)

    install.cli(subparser)
    images.cli(subparser)
    mount.cli(subparser)
    containers.cli(subparser)
    push.cli(subparser)
    pull.cli(subparser)
    run.cli(subparser)
    scan.cli(subparser)
    stop.cli(subparser)
    storage.cli(subparser)
    top.cli(subparser)

    # atomic rhost (if commctl.cli is available)
    try:
        import commctl.cli
        rhostp = subparser.add_parser(
            "rhost", help=_("execute cluster host management commands"))
        commctl.cli.add_host_commands(rhostp)
    except ImportError:
        pass

    uninstall.cli(subparser)

    mount.cli_unmount(subparser)
    update.cli(subparser)
    verify.cli(subparser)
    info.cli_version(subparser)

    return parser

if __name__ == '__main__':
    try:
        with Atomic.Atomic() as atomic:
            aparser = create_parser(atomic.help())
            args = aparser.parse_args()
            _class = atomic if '_class' not in args else args._class() # pylint: disable=protected-access
            _class.set_args(args)
            _func = getattr(_class, args.func)
            sys.exit(_func())
    except KeyboardInterrupt:
        sys.exit(0)
    except (ValueError, IOError, docker.errors.DockerException, NoDockerDaemon) as e:
        write_err("%s" % str(e))
        if os.geteuid() != 0:
            need_root()
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        write_err("")
        sys.exit(e.returncode)
    except AttributeError:
        # python3 throws exception on no args to atomic
        aparser.print_usage()
        sys.exit(1)
    except MountError as e:
        if str(e).find("Permission denied") > 0:
            need_root()
        else:
            write_err("%s" % str(e))
        sys.exit(1)
    except Exception as e: # pylint: disable=broad-except
        write_err("%s" % str(e))
        sys.exit(1)
    except SystemExit as e:
        # Overriding debug args to avoid a traceback from sys.exit()
        if 'args' in locals():
            args.debug = False
        sys.exit(e.code)
    finally:
        if 'args' in locals() and args.debug:
            traceback.print_exc(file=sys.stdout)