#!/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.diff import Diff
from Atomic.top import Top
from Atomic.verify import Verify

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
    builtins.__dict__['_'] = str


class HelpByDefaultArgumentParser(argparse.ArgumentParser):

    def error(self, message):
        sys.stderr.write('%s: %s\n' % (sys.argv[0], message))
        sys.stderr.write("Try '%s --help' for more information.\n" % self.prog)
        sys.exit(2)


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

    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')
            self._choices_actions.pop()
            pseudo_action = self._AliasedPseudoAction(name, aliases, help)
            self._choices_actions.append(pseudo_action)

        return parser

def add_opt(sub):
    sub.add_argument("--opt1", dest="opt1",help=argparse.SUPPRESS)
    sub.add_argument("--opt2", dest="opt2",help=argparse.SUPPRESS)
    sub.add_argument("--opt3", dest="opt3",help=argparse.SUPPRESS)


def check_negative(value):
    ivalue = int(value)
    if ivalue < 1:
        raise argparse.ArgumentTypeError("%s must be a positive integer value greater than 0." % value)
    return ivalue


if __name__ == '__main__':
    if os.geteuid() != 0:
        exit("%s needs to be run as root." % sys.argv[0] )
   
    atomic = Atomic.Atomic()
    parser = HelpByDefaultArgumentParser(description=atomic.help())
    parser.register('action', 'parsers', AliasedSubParsersAction)
    parser.add_argument('-v', '--version', action='version', version=Atomic.__version__)
    subparser = parser.add_subparsers(help=_("commands"))

    # atomic diff
    diffp = subparser.add_parser(
    "diff", help=_("Show differences between two container images, file diff or RPMS."),
    epilog="atomic diff 'image1|container1' 'image2|container2'")
    diffp.set_defaults(_class=Diff, func='diff')
    diffp.add_argument("compares", nargs=2,
                       help=_("Container images to compare"))
    diffp.add_argument("--json", default=False, action='store_true',
                       help=_("output json"))
    diffp.add_argument("-n", "--no-files", default=False, action='store_true',
                       help=_("Do not perform a file diff between the docker objects"))
    diffp.add_argument("--names-only", default=False,
                       action='store_true', help=_("Only show RPM names and not versions"))
    diffp.add_argument("-r", "--rpms", default=False, action='store_true',
                       help=_("List different rpms between the container images."))
    diffp.add_argument("-v", "--verbose", default=False, action='store_true',
                       help=_("Show verbose output, listing all RPMs"))

    if os.path.exists("/usr/bin/rpm-ostree"):
        # atomic host
        hostp = subparser.add_parser("host", help=_("execute Atomic host "
                                                    "commands"))
        host_subparser = hostp.add_subparsers(help=_("host commands"))

        # atomic host rollback
        rollbackp = host_subparser.add_parser(
            "rollback", help=_("switch to alternate installed tree at next boot"))
        rollbackp.set_defaults(func='host_rollback')
        rollbackp.add_argument("-r", "--reboot", dest="reboot",
                               action="store_true",
                               help=_("initiate a reboot after rollback is "
                                      "prepared"))
        rollbackp.add_argument("args", nargs=argparse.REMAINDER,
                               help=_("Additional arguments appended to the "
                                      "rollback method.  Use `-- --OPTION=VAL` "
                                      "if you want to pass additional "
                                      "unsupported arguments to rpm-ostree."))


        # atomic host status
        statusp = host_subparser.add_parser(
            "status", help=_("list information about all deployments"))
        statusp.add_argument("-p", "--pretty", dest="pretty",
                              action="store_true",
                              help=_("Display status in formatted rows"))
        statusp.add_argument("args", nargs=argparse.REMAINDER,
                             help=_("Additional arguments appended to the "
                                    "status method.  Use `-- --OPTION=VAL` "
                                    "if you want to pass additional "
                                    "unsupported arguments to rpm-ostree."))

        statusp.set_defaults(func='host_status')

        # atomic host upgrade
        upgradep = host_subparser.add_parser(
            "upgrade", help=_("upgrade to the latest Atomic tree if one "
                              "is available"))
        upgradep.set_defaults(func='host_upgrade')
        upgradep.add_argument("-r", "--reboot", dest="reboot",
                              action="store_true",
                              help=_("if an upgrade is available, reboot "
                                     "after deployment is complete"))
        upgradep.add_argument("--allow-downgrade", dest="downgrade",
                              action="store_true",
                              help=_("Permit deployment of chronologically older trees"))
        upgradep.add_argument("--os", dest="os",
                              help=_("Operate on provided OSNAME"))
        upgradep.add_argument("--check-diff", dest="diff",
                              action="store_true",
                              help=_("Check for upgrades and print package diff only"))
        upgradep.add_argument("args", nargs=argparse.REMAINDER,
                              help=_("Additional arguments appended to the "
                                     "upgrade method.  Use `-- --OPTION=VAL` "
                                     "if you want to pass additional "
                                     "unsupported arguments to rpm-ostree."))

        # atomic host rebase
        rebasep = host_subparser.add_parser(
            "rebase", help=_("Download and deploy a new origin refspec"))
        rebasep.set_defaults(func='host_rebase')
        rebasep.add_argument("--os", dest="os",
                              help=_("Operate on provided OSNAME"))
        rebasep.add_argument("refspec",
                             help=_("Origin refspec for new deployment"))
        rebasep.add_argument("args", nargs=argparse.REMAINDER,
                             help=_("Additional arguments appended to the "
                                    "rebase method.  Use `-- --OPTION=VAL` "
                                    "if you want to pass additional "
                                    "unsupported arguments to rpm-ostree."))

        # atomic host deploy
        deployp = host_subparser.add_parser(
            "deploy", help=_("deploy a specific commit"))
        deployp.set_defaults(func='host_deploy')
        deployp.add_argument("revision", help=_("Checksum or version to deploy"));
        deployp.add_argument("-r", "--reboot", dest="reboot",
                             action="store_true",
                             help=_("Reboot after deployment is complete"))
        deployp.add_argument("--os", dest="os",
                             help=_("Operate on provided OSNAME"))
        deployp.add_argument("--preview", dest="preview",
                             action="store_true",
                             help=_("Just preview package differences"))
        deployp.add_argument("args", nargs=argparse.REMAINDER,
                             help=_("Additional arguments appended to the "
                                    "deploy method.  Use `-- --OPTION=VAL` "
                                    "if you want to pass additional "
                                    "unsupported arguments to rpm-ostree."))

    # atomic info
    infop = subparser.add_parser(
        "info", help=_("display label information about an image"),
        epilog="atomic info attempts to read and display the LABEL "
        "information about an image")
    infop.set_defaults(func='info')
    infop.add_argument("--remote", dest="force_remote_info",
                       action='store_true', default=False,
                       help=_('ignore local images and only scan registries'))
    infop.add_argument("image", help=_("container image"))

    # atomic install
    installp = subparser.add_parser(
        "install", help=_("execute container image install method"),
        epilog="atomic install attempts to read the LABEL INSTALL field "
        "in the image, if it does not exist atomic will just pull "
        "the image on to your machine.  You could add a LABEL "
        "INSTALL command to your Dockerfile like: 'LABEL INSTALL "
        "%s'" % atomic.print_install())
    installp.set_defaults(func='install')
    add_opt(installp)
    installp.add_argument("-n", "--name", dest="name", default=None,
                          help=_("name of container"))
    installp.add_argument(
        "--display",
        default=False,
        action="store_true",
        help=_("preview the command that %s would execute") % sys.argv[0])
    installp.add_argument("image", help=_("container image"))
    installp.add_argument("args", nargs=argparse.REMAINDER,
                          help=_("Additional arguments appended to the image "
                                 "install method"))

    # atomic images
    imagesp = subparser.add_parser(
        "images", help=_("list container images on your system"),
        epilog="atomic images by default will list all installed "
        "container images on your system.  Using the --prune "
        "option, will free up disk space deleting unused "
        "'dangling' images")
    imagesp.set_defaults(func='images')
    imagesp.add_argument("--prune", action="store_true",
                         help=_("prune dangling images"))

    # atomic migrate
    migratep = subparser.add_parser("migrate", help=_("migrate container's content from one backend storage "
                                                    "to another"))
    migrate_subparser = migratep.add_subparsers(help=_("migrate commands"))

   # atomic export
    exportp = migrate_subparser.add_parser(
        "export", help=_("export containers and it's associated contents into a filesystem directory"),
        epilog="export container's content. "
        "The export command would export images, "
        "containers and volumes into a filesystem directory.")
    exportp.set_defaults(func='Export')
    exportp.add_argument("--graph", dest="graph", default="/var/lib/docker",
                        help=_("Root of the Docker runtime (Default: /var/lib/docker)"))
    exportp.add_argument("--dir", dest="export_location",
                         default="/var/lib/atomic/migrate",
                         help=_("Path for exporting container's content (Default: /var/lib/atomic/migrate)"))
	
    # atomic import
    importp = migrate_subparser.add_parser(
         "import", help=_("import containers and it's associated contents from a filesystem directory"),
         epilog="import container's content."
         "The import command would import images,"
         "containers and volumes from a filesystem directory.")
    importp.set_defaults(func='Import')
    importp.add_argument("--graph", dest="graph", default="/var/lib/docker",
                         help=_("Root of the Docker runtime (Default: /var/lib/docker)"))
    importp.add_argument("--dir", dest="import_location",
                         default="/var/lib/atomic/migrate",
                         help=_("Path for importing container's content (Default: /var/lib/atomic/migrate)"))

    # atomic mount
    mountp = subparser.add_parser(
        "mount", help=_("mount container image to a specified directory"),
        epilog="atomic mount attempts to mount a container image to a "
        "specified directory so that its contents may be "
        "inspected.")
    mountp.set_defaults(func='mount')
    mountp.add_argument("-o", "--options", dest="options", default="",
                        help=_("comma-separated list of mount options, "
                               "defaults are 'ro,nodev,nosuid'"))
    mountp.add_argument("--live", dest="live", action="store_true",
                        help=_("mount a running container 'live', allowing "
                               "modification of the contents."))
    mountp.add_argument("image", help=_("image/container id"))
    mountp.add_argument("mountpoint", help=_("filesystem location to mount "
                                             "the image/container"))
    # atomic push
    pushp = subparser.add_parser(
        "push", aliases=['upload'], help=_("push latest image to repository"),
        epilog="push the latest specified image to a repository.")
    pushp.set_defaults(func='push')

    # making it so we cannot call both the --pulp and --satellite commands
    # at the same time (mutually exclusive)
    pushgroup = pushp.add_mutually_exclusive_group()
    pushgroup.add_argument("--pulp",
                           default=False,
                           action="store_true",
                           help=_("push image using pulp"))
    pushgroup.add_argument("--satellite",
                           default=False,
                           action="store_true",
                           help=_("push image using Satellite"))

    pushp.add_argument("--verify_ssl",
                         default=None,
                         action="store_true",
                         help=_("flag to verify ssl of registry"))
    pushp.add_argument("--debug",
                         default=None,
                         action="store_true",
                         help=_("debug mode"))
    pushp.add_argument("-U", "--url",
                         dest="url",
                         default=None,
                         help=_("URL for remote registry"))
    pushp.add_argument("-u", "--username",
                         default=None,
                         dest="username",
                         help=_("Username for remote registry"))
    pushp.add_argument("-p", "--password",
                         default=None,
                         dest="password",
                         help=_("Password for remote registry"))
    pushp.add_argument("image", help=_("container image"))
    pushp.add_argument("-a", "--activation_key",
                         default=None,
                         dest="activation_key",
                         help=_("Activation Key"))
    pushp.add_argument("-r", "--repository_id",
                         default=None,
                         dest="repo_id",
                         help=_("Repository ID"))
    # pushp.add_argument("--activation_key_name",
    #                      default=None,
    #                      dest="activation_key_name",
    #                      help=_("Activation Key Name"))
    # pushp.add_argument("--repo_name", "--repository_name",
    #                      default=None,
    #                      dest="repo_name",
    #                      help=_("Repository Name"))
    # pushp.add_argument("--org_name", "--organization_name",
    #                      default=None,
    #                      dest="org_name",
    #                      help=_("Organization Name"))

    # atomic scan
    def bool_detect(myarg):
        return True if myarg.lower() in \
            ('yes', 'true', 't', 'y', '1') else False

    scanp = subparser.add_parser(
        "scan", help=_("scan an image or container for CVEs"),
        epilog="atomic scan <input> scans a container or image for CVEs")
    scanp.set_defaults(func='scan')
    scan_group = scanp.add_mutually_exclusive_group()
    scan_out = scanp.add_mutually_exclusive_group()
    scan_out.add_argument("--json", default=False, action='store_true', help=_("output json"))
    scan_out.add_argument("--detail", default=False, action='store_true', help=_("output more detail"))
    scanp.add_argument("scan_targets", nargs='*', help=_("container image"))
    scanp.add_argument("--fetch_cves", type=bool_detect, default=None,  help=_("override the behavior defined in /etc/oscapd/config.ini as to whether to download the most recent CVE data. Values can be True of False"))
    scan_group.add_argument("--all", default=False, action='store_true', help=_("scan all images (excluding intermediate layers) and containers"))
    scan_group.add_argument("--images", default=False, action='store_true', help=_("scan all images (excluding intermediate layers)"))
    scan_group.add_argument("--containers", default=False, action='store_true', help=_("scan all containers"))

    # atomic stop
    stopp = subparser.add_parser(
        "stop", help=_("execute container image stop method"),
        epilog="atomic will just stop the container if it is running, if "
        "image does not specify LABEL STOP")
    stopp.set_defaults(func='stop')
    add_opt(stopp)
    stopp.add_argument("-n", "--name", dest="name", default=None,
                       help=_("name of container"))
    stopp.add_argument("image", help=_("container image"))
    stopp.add_argument("args", nargs=argparse.REMAINDER,
                          help=_("Additional arguments appended to the image "
                                 "stop method"))

    # atomic run
    runp = subparser.add_parser(
        "run", help=_("execute container image run method"),
        epilog="atomic run defaults to the following command, if image "
        "does not specify LABEL run\n'%s'" % atomic.print_run())
    runp.set_defaults(func='run')
    add_opt(runp)
    runp.add_argument("-n", "--name", dest="name", default=None,
                      help=_("name of container"))
    runp.add_argument("--spc", default=False, action="store_true",
                      help=_("use super privileged container mode: '%s'" %
                             atomic.print_spc()))
    runp.add_argument("image", help=_("container image"))
    runp.add_argument("command", nargs=argparse.REMAINDER,
                      help=_("command to execute within the container. "
                             "If container is not running, command is appended"
                             "to the image run method"))

    runp.add_argument(
        "--display",
        default=False,
        action="store_true",
        help=_("preview the command that %s would execute") % sys.argv[0])

    # atomic top
    topp = subparser.add_parser(
    "top", help=_("Show top-like stastics about processes running in containers"))
    topp.set_defaults(_class=Top, func='atomic_top')
    topp.add_argument("-d", type=int, default=1, help=_("Interval (secs) to refresh process information"))
    topp.add_argument("-o", "--optional", help=_("Additional fields to display"), nargs='*', choices=['time', 'stime', 'ppid', 'uid', 'gid', 'user', 'group'])
    topp.add_argument("-n", help=_("Number of iterations"), type=check_negative)
    topp.add_argument("containers", nargs="*", help=_("list of containers to monitor, leave blank for all"))

    # atomic uninstall
    uninstallp = subparser.add_parser(
        "uninstall", help=_("execute container image uninstall method"),
        epilog="atomic uninstall attempts to read the LABEL UNINSTALL "
        "field in the image, if it does not exist atomic will "
        "remove the image from your machine.  You could add a "
        "LABEL UNINSTALL command to your Dockerfile like: 'LABEL "
        "UNINSTALL %s'" % atomic.print_uninstall())
    uninstallp.set_defaults(func='uninstall')
    add_opt(uninstallp)
    uninstallp.add_argument("-n", "--name", dest="name", default=None,
                            help=_("name of container"))
    uninstallp.add_argument("-f", "--force", default=False, dest="force",
                            action="store_true",
                            help=_("remove all containers based on this "
                                   "image"))
    uninstallp.add_argument("image", help=_("container image"))
    uninstallp.add_argument("args", nargs=argparse.REMAINDER,
                            help=_("Additional arguments appended to the "
                                   "image uninstall method"))

    # atomic unmount
    unmountp = subparser.add_parser(
        "unmount", aliases=["umount"],help=_("unmount container image"),
        epilog="atomic unmount will unmount a container image previously "
        "mounted with atomic mount")
    unmountp.set_defaults(func='unmount')
    unmountp.add_argument("mountpoint",
                          help=_("filesystem location of image/container to "
                                 "be unmounted"))

    # atomic update
    updatep = subparser.add_parser(
        "update", help=_("pull latest container image from repository"),
        epilog="atomic update downloads the latest container image. If a "
        "previously created  container based on this image exists, "
        "the container will continue to use the old image.  Use "
        "--force to remove the outdated container.")
    updatep.set_defaults(func='update')
    updatep.add_argument("-f", "--force", default=False, dest="force",
                         action="store_true",
                         help=_("remove all containers based on this image"))
    updatep.add_argument("image", help=_("container image"))

    # atomic version
    versionp = subparser.add_parser(
        "version", help=_("display image 'Name Version Release' label"),
        epilog="atomic version displays the image version information, if "
        "it is provided")
    versionp.add_argument("-r", "--recurse", default=False, dest="recurse",
                          action="store_true",
                          help=_("recurse through all layers"))
    versionp.set_defaults(func='print_version')
    versionp.add_argument("image", help=_("container image"))

    # atomic verify
    verifyp = subparser.add_parser(
        "verify", help=_("verify image is fully updated"),
        epilog="atomic verify checks whether there is a newer image "
        "available and scans through all layers to see if any of "
        "the sublayers have a new version available")
    verifyp.set_defaults(_class=Verify, func='verify')
    verifyp.add_argument("image", help=_("container image"))
    verifyp.add_argument("-v", "--verbose", default=False,
                          action="store_true",
                          help=_("Report status of each layer"))

    try:
        args = parser.parse_args()
        _class = atomic if '_class' not in args else args._class()
        _class.set_args(args)
        _func = getattr(_class, args.func)
        sys.exit(_func())
    except KeyboardInterrupt:
        sys.exit(0)
    except (ValueError, IOError, docker.errors.DockerException) as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        sys.stderr.write("\n")
        sys.exit(e.returncode)
    except AttributeError:
        # python3 throws exception on no args to atomic
        parser.print_usage()
        sys.exit(1)
