#!/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 sys, os
import argparse
import gettext
import docker
import json
import subprocess
import requests

PROGNAME="atomic"
gettext.bindtextdomain(PROGNAME, "/usr/share/locale")
gettext.textdomain(PROGNAME)
try:
    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 Atomic(object):
    INSTALL_ARGS = ["/usr/bin/docker", "run",
                    "-t",
                    "-i",
                    "--rm",
                    "--privileged",
                    "-v", "/:/host",
                    "--net=host",
                    "--ipc=host",
                    "--pid=host",
                    "-e", "HOST=/host",
                    "-e", "NAME=NAME",
                    "-e", "IMAGE=IMAGE",
                    "-v", "${CONFDIR}:/etc/NAME",
                    "-v", "${LOGDIR}:/var/log/NAME",
                    "-v", "${DATADIR}:/var/lib/NAME",
                    "-e", "CONFDIR=${CONFDIR}",
                    "-e", "LOGDIR=${LOGDIR}",
                    "-e", "DATADIR=${DATADIR}",
                    "--name", "NAME",
                    "IMAGE"]

    SPC_ARGS = ["/usr/bin/docker", "run",
                "-t",
                "-i",
                "--rm",
                "--privileged",
                "-v", "/:/host",
                "-v", "/run:/run",
                "-v", "/etc/localtime:/etc/localtime",
                "--net=host",
                "--ipc=host",
                "--pid=host",
                "-e", "HOST=/host",
                "-e", "NAME=NAME",
                "-e", "IMAGE=IMAGE",
                "IMAGE" ]

    RUN_ARGS = ["/usr/bin/docker", "create",
                "-t",
                "-i",
                "--name", "NAME",
                "IMAGE" ]

    def __init__(self):
        self.d = docker.Client()
        self.name = None
        self.image = None
        self.spc = False
        self.inspect = None
        self.force = False

    def writeOut(self, output, lf="\n"):
        sys.stdout.flush()
        sys.stdout.write(output + lf)

    def force_delete_containers(self):
        if self._inspect_image():
            image = self.image
            if self.image.find(":") == -1:
                image += ":latest"
            for c in self.d.containers(all=True):
                if c["Image"] == image:
                    self.d.remove_container(c["Id"], force=True)

    def update(self):
        if self.force:
            self.force_delete_containers()
        return subprocess.check_call(["/usr/bin/docker", "pull", self.image])

    def pull(self):
        prevstatus = ""
        prev = ""
        for line in self.d.pull(self.image, stream=True):
            bar = json.loads(line)
            status = bar['status']
            if prevstatus != status:
                self.writeOut(status, "")
            if 'id' not in bar:
                continue
            if status == "Downloading":
                self.writeOut(bar['progress'] + " ")
            elif status == "Extracting":
                self.writeOut("Extracting: " + bar['id'])
            elif status == "Pull complete":
                pass
            elif status.startswith("Pulling"):
                self.writeOut("Pulling: " + bar['id'])

            prevstatus = status
        self.writeOut("")

    def set_args(self, args):
        self.args=args
        try:
            self.image = args.image
        except:
            pass
        try:
            self.command = args.command
        except:
            self.command = None

        try:
            self.spc = args.spc
        except:
            self.spc = False

        try:
            self.name = args.name
        except:
            pass

        try:
            self.force = args.force
        except:
            pass

        if not self.name and self.image is not None:
            self.name = self.image.split("/")[-1].split(":")[0]
            if self.spc:
                self.name = self.name + "-spc"


    def _getconfig(self, key, default=None):
        assert self.inspect is not None
        cfg = self.inspect.get("Config")
        if cfg is None:
            return default
        val = cfg.get(key, default)
        if val is None:
            return default
        return val

    def _get_cmd(self):
        return self._getconfig("Cmd", [ "/bin/sh" ])

    def _get_labels(self):
        return self._getconfig("Labels", [])

    def _interactive(self):
        return (self._getconfig("AttachStdin", False) and
                self._getconfig("AttachStdout", False) and
                self._getconfig("AttachStderr", False))

    def _running(self):
        if self._interactive():
            cmd = ["/usr/bin/docker", "exec", "-t", "-i", self.name]
            if self.command:
                cmd += self.command
            else:
                cmd += self._get_cmd()
            return subprocess.check_call(cmd, stderr=subprocess.PIPE)
        else:
            self.writeOut("Container is running")

    def _start(self):
        if self._interactive():
            if self.command:
                subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)
                return subprocess.check_call(["/usr/bin/docker", "exec", "-t", "-i", self.name] + self.command)
            else:
                return subprocess.check_call(["/usr/bin/docker", "start", "-i", "-a", self.name], stderr=subprocess.PIPE)
        else:
            if self.command:
                subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)
                return subprocess.check_call(["/usr/bin/docker", "exec", "-t", "-i", self.name] + self.command)
            else:
                return subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)

    def _inspect_image(self):
        try:
            return self.d.inspect_image(self.image)
        except docker.errors.APIError:
            pass
        except requests.exceptions.ConnectionError as e:
            raise IOError("Unable to communicate with docker daemon: %s\n" % str(e))
        return None

    def _inspect_container(self):
        try:
            return self.d.inspect_container(self.name)
        except docker.errors.APIError:
            pass
        except requests.exceptions.ConnectionError as e:
            raise IOError("Unable to communicate with docker daemon: %s\n" % str(e))
        return None

    def _get_args(self, label):
        labels = self._get_labels()
        if label in labels:
            return labels[label].split()
        return None

    def _check_latest(self):
        inspect = self._inspect_image()
        if inspect["Id"] != self.inspect["Image"]:
            response = ""
            sys.stdout.write("""The '%(name)s' container is using an older version of the installed
'%(image)s' container image. If you wish to use the newer image,
you must either create a new container with a new name or
uninstall the '%(name)s' container.

# atomic uninstall --name %(name)s %(image)s

and create new container on the '%(image)s' image.

atomic update --force %(image)s

removes all containers based on an image.

""" % { "name" : self.name, "image" : self.image} )

    def container_run_command(self):
        command = "%s run " % sys.argv[0]
        if self.spc:
            command += "--SPC "

        if self.name != self.image:
            command += "--name %s " % self.name
        command += self.image
        return command

    def run(self):
        self.inspect = self._inspect_container()

        if self.inspect:
            self._check_latest()
            # Container exists
            if self.inspect["State"]["Running"]:
                return self._running()
            else:
                return self._start()
        else:
            if self.command:
                raise ValueError("Container '%s' must be running before executing a command into it.\nExecute the following to create the container:\n%s" % (self.name, self.container_run_command()))

        # Container does not exist
        self.inspect = self._inspect_image()
        if not self.inspect:
            self.update()
            self.inspect = self._inspect_image()

        if self.spc:
            if self.command:
                args = self.SPC_ARGS + self.command
            else:
                args = self.SPC_ARGS +  self._get_cmd()

            cmd = self.gen_cmd(args)
            self.writeOut(cmd)
        else:
            missing_RUN = False
            args = self._get_args("RUN")
            if not args:
                missing_RUN = True
                args = self.RUN_ARGS + self._get_cmd()

            cmd = self.gen_cmd(args)
            self.writeOut(cmd)

            if missing_RUN:
                subprocess.check_call(cmd, env={
                    "CONFDIR": "/etc/%s" % self.name,
                    "LOGDIR": "/var/log/%s" % self.name,
                    "DATADIR":"/var/lib/%s" % self.name}, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
                return self._start()

        subprocess.check_call(cmd, env={
            "CONFDIR": "/etc/%s" % self.name,
            "LOGDIR": "/var/log/%s" % self.name,
            "DATADIR":"/var/lib/%s" % self.name}, shell=True)


    def stop(self):
        self.inspect = self._inspect_container()
        args = self._get_args("STOP")
        if args:
            cmd = self.gen_cmd(args)
            self.writeOut(cmd)

            subprocess.check_call(cmd, env={
                "CONFDIR": "/etc/%s" % self.name,
                "LOGDIR": "/var/log/%s" % self.name,
                "DATADIR":"/var/lib/%s" % self.name}, shell=True)


        # Container exists
        if self.inspect["State"]["Running"]:
            self.d.stop(self.name)

    def _rpmostree(self, *args):
        os.execl("/usr/bin/rpm-ostree", "rpm-ostree", *args)

    def host_status(self):
        self._rpmostree("status")

    def host_upgrade(self):
        argv = ["upgrade"]
        if self.args.reboot:
            argv.append("--reboot")
        self._rpmostree(*argv)

    def host_rollback(self):
        argv = ["rollback"]
        if self.args.reboot:
            argv.append("--reboot")
        self._rpmostree(*argv)

    def uninstall(self):
        self.inspect = self._inspect_container()

        if self.force:
            self.force_delete_containers()

        if self.name != self.image:
            try:
                # Attempt to remove container, if it exists just return
                self.d.stop(self.name)
                self.d.remove_container(self.name)
                return
            except:
                # On exception attempt to remove image
                pass

        try:
            self.d.stop(self.image)
            self.d.remove_container(self.image)
        except docker.errors.APIError:
            pass
        self.inspect = self._inspect_image()
        args = self._get_args("UNINSTALL")
        if args:
            cmd = self.gen_cmd(args)
            self.writeOut(cmd)
            subprocess.check_call(cmd, env={
                "CONFDIR": "/etc/%s" % self.name,
                "LOGDIR": "/var/log/%s" % self.name,
                "DATADIR":"/var/lib/%s" % self.name}, shell=True)
        subprocess.check_call(["/usr/bin/docker", "rmi", self.image])

    def gen_cmd(self,cargs):
        args = []
        for c in cargs:
            if c == "IMAGE":
                args.append(self.image)
                continue
            if c == "IMAGE=IMAGE":
                args.append("IMAGE=%s" % self.image)
                continue
            if c == "NAME=NAME":
                args.append("NAME=%s" % self.name)
                continue
            if c == "NAME":
                args.append(self.name)
                continue
            args.append(c)
        return " ".join(args)

    def info(self):
        self.inspect = self._inspect_image()
        if not self.inspect:
            self.update()
            self.inspect = self._inspect_image()

        labels = self._get_labels()
        for label in labels:
            self.writeOut("%-13s: %s" % (label, labels[label]))

    def install(self):
        self.inspect = self._inspect_image()
        if not self.inspect:
            self.update()
            self.inspect = self._inspect_image()

        args = self._get_args("INSTALL")
        if args:
            cmd = self.gen_cmd(args)
            self.writeOut(cmd)

            return(subprocess.check_call(cmd, env={
                "CONFDIR": "/etc/%s" % self.name,
                "LOGDIR": "/var/log/%s" % self.name,
                "DATADIR":"/var/lib/%s" % self.name}, shell=True))

    def help(self):
        if os.path.exists("/usr/bin/rpm-ostree"):
            return _('Atomic Management Tool')
        else:
            return _('Atomic Container Tool')

    def print_spc(self):
        return " ".join(self.SPC_ARGS)

    def print_run(self):
        return " ".join(self.RUN_ARGS)

    def print_install(self):
        return " ".join(self.INSTALL_ARGS) + " /usr/bin/INSTALLCMD"

    def print_uninstall(self):
        return " ".join(self.INSTALL_ARGS) + " /usr/bin/UNINSTALLCMD"

def SetFunc(function):
    class customAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            setattr(namespace, self.dest, function)
    return customAction

class HelpByDefaultArgumentParser(argparse.ArgumentParser):
    def error(self, message):
        self.print_help()
        sys.stderr.write('\nerror: %s\n' % message)
        sys.exit(2)

if __name__ == '__main__':
    atomic=Atomic()
    parser = HelpByDefaultArgumentParser(description=atomic.help())
    subparser = parser.add_subparsers(help=_("commands"))

    if os.path.exists("/usr/bin/rpm-ostree"):
        hostp = subparser.add_parser("host",help=_("execute Atomic host commands"))
        host_subparser = hostp.add_subparsers(help=_("host commands"))
        rollbackp = host_subparser.add_parser("rollback", help=_("switch to alternate tree at boot"))
        rollbackp.set_defaults(func=atomic.host_rollback)
        rollbackp.add_argument("-r", "--reboot", dest="reboot",
                               action="store_true",
                               help=_("initiate a reboot after rollback is prepared"))

        statusp = host_subparser.add_parser("status", help=_("list information about all deployments"))
        statusp.set_defaults(func=atomic.host_status)
        upgradep = host_subparser.add_parser("upgrade", help=_("upgrade to the latest Atomic tree if one is available"))
        upgradep.set_defaults(func=atomic.host_upgrade)
        upgradep.add_argument("-r", "--reboot", dest="reboot",
                              action="store_true",
                              help=_("if an upgrade is available, reboot after deployment is complete"))

    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=atomic.info)
    infop.add_argument("image", help=_("container image"))

    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=atomic.install)
    installp.add_argument("image", help=_("container image"))
    installp.add_argument("command", nargs=argparse.REMAINDER,
                          help=_("execute container image install method"))
    installp.add_argument("-n", "--name", dest="name",
                      default=None,
                      help=_("name of container"))
    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\n")
    stopp.set_defaults(func=atomic.stop)
    stopp.add_argument("-n", "--name", dest="name",
                      default=None,
                      help=_("name of container"))
    stopp.add_argument("image", help=_("container image"))

    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=atomic.run)
    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"))

    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=atomic.uninstall)
    uninstallp.add_argument("-n", "--name", dest="name",
                            default=None,
                            help=_("name of container"))
    uninstallp.add_argument("-f", "--force",
                            default=False,
                            action="store_true",
                            help=_("remove all containers based on this image"))
    uninstallp.add_argument("image", help=_("container image"))

    updatep = subparser.add_parser("update",help=_("pull latest container image from repository"), epilog="atomic update downloads the latest container image. If a previously container based on this image exists, the container will continue to use the old image.  Use --force to remove the container.")

    updatep.set_defaults(func=atomic.update)
    updatep.add_argument("-f", "--force",
                         default=False,
                         action="store_true",
                         help=_("remove all containers based on this image"))
    updatep.add_argument("image", help=_("container image"))

    try:
        args = parser.parse_args()
        atomic.set_args(args)
        sys.exit(args.func())
    except ValueError as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
    except IOError as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(0)
    except subprocess.CalledProcessError as e:
        sys.stderr.write("\n")
        sys.exit(e.returncode)
    except docker.errors.DockerException as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
