#!/bin/sh

''':'
PYTHON=$(head -n 1 "$GPHOME/lib/python/gppylib/gplog.py" 2> /dev/null | cut -d ' ' -f 2)

if [ -z $PYTHON ]; then
    PYTHON=$(which python3 2> /dev/null || which python 2> /dev/null || which python2 2> /dev/null)
fi

if [ -z $PYTHON ]; then
    echo "no python3 or python2 detected"
    exit 1
fi

exec $PYTHON $0 "$@"
exit $?
''' #'''

#------------------------------------------------------------------------------
#
# Copyright (c) 2016-Present VMware, Inc. or its affiliates.
#
#------------------------------------------------------------------------------

import os
import re
import sys
import hashlib
import tempfile
import traceback
import datetime
import xml.dom.minidom
import subprocess
import argparse
from xml.etree import ElementTree

try:
    from gppylib import gplog
    from gppylib.commands.base import WorkerPool, REMOTE, Command
    from gppylib.commands.unix import getLocalHostname, getUserName
    try:
        # https://github.com/greenplum-db/gpdb/pull/14145
        # after gp7 the Scp tool renamed to Rsync
        from gppylib.commands.unix import Scp as Rsync
    except ImportError:
        from gppylib.commands.unix import Rsync
    from gppylib.db import dbconn
    from gppylib import gparray
    from gppylib import userinput
except ImportError as e:
    sys.exit('ERROR: Cannot import Greenplum modules.  Please check that you have sourced greenplum_path.sh.  Detail: ' + str(e))

PLCONTAINER_CONFIG_FILE = 'plcontainer_configuration.xml'
PLCONTAINER_GPDB_DIR = 'share/postgresql/plcontainer'
PLCONTAINER_CLIENT_DIR = "bin/plcontainer_clients"

PL_BATCH_SIZE = 64
RUNTIME_ID_MAX_LENGTH = 64
TMP = '/tmp'
TAR_GZ = '.tar.gz'

global DBHOME

def parseargs():
    parser = argparse.ArgumentParser()

    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    subparsers = parser.add_subparsers(dest='subparser_name')

    image_add_parser = subparsers.add_parser("image-add", help = "Install specified docker image on all hosts")
    image_delete_parser = subparsers.add_parser("image-delete", help = "Delete specified docker image on all hosts")
    image_list_parser = subparsers.add_parser("image-list", help = "Show existing container images on the current host")

    runtime_add_parser = subparsers.add_parser("runtime-add", help = "Add a container runtime configuration on all hosts")
    runtime_replace_parser = subparsers.add_parser("runtime-replace", help = "Replace with a container runtime configuration on all hosts")
    runtime_show_parser = subparsers.add_parser("runtime-show", help = "Show container runtime configurations on the current host")
    runtime_delete_parser = subparsers.add_parser("runtime-delete", help = "Delete a given container runtime configuration on all hosts")
    runtime_edit_parser = subparsers.add_parser("runtime-edit", help = "Use editor (vi by default) to edit and then update configurations on all hosts")
    runtime_backup_parser = subparsers.add_parser("runtime-backup", help = "Backup the local runtime configurations to a file")
    runtime_restore_parser = subparsers.add_parser("runtime-restore", help = "Restore the runtime configurations from a file and distribute them to all hosts")
    runtime_verify_parser = subparsers.add_parser("runtime-verify",
                                                help="Sanity check configurations on all the hosts")

    # For image-add
    image_add_options = image_add_parser.add_mutually_exclusive_group(required=True)
    image_add_options.add_argument("-f", "--file", dest="file", help="image location (e.g. /tmp/centos.7.2.tar.gz)")
    image_add_parser.add_argument("-ulc", "--use_local_copy", action="store_true", dest="local", help="If set, use the local image file (default is no)")
    image_add_options.add_argument("-u", "--url", dest="url", help="image url (e.g. 192.168.0.1:5000/images/centos:7.2)")

    # For image-delete
    image_delete_parser.add_argument("-i", "--image", dest="image_name", help="image name:tag (e.g. centos:7.2)", required=True)

    # For runtime-add
    runtime_add_parser.add_argument("-r", "--runtime", dest="runtime_id",
                                    help="unique string which presents configurations for a container runtime)", required=True)
    runtime_add_parser.add_argument("-i", "--image", dest="image_name",
                                    help="image name:tag (e.g. centos:7.2)", required=True)
    runtime_add_parser.add_argument("-l", "--language", dest="language",
                                    help="language type, currently support python, python3 and r. " +
                                    "Used to auto-configure container command.", required=True, choices=['r','python','python3'])
    runtime_add_parser.add_argument("-v", "--volume", dest="shared", action="append",
                                    help="Bind mount a volume, can work with multiple volumes" +
                                    " format: HOST:CONTAINER:[rw|ro]")
    runtime_add_parser.add_argument("-s", "--setting", dest="setting", action="append",
                                    help="Setting for the runtime. Allow call multiple times." +
                                    " (e.g. memory_mb=256). See more in document.")

    # For runtime-replace
    runtime_replace_parser.add_argument("-r", "--runtime", dest="runtime_id",
                                    help="unique string which presents configurations for a container runtime)", required=True)
    runtime_replace_parser.add_argument("-i", "--image", dest="image_name",
                                    help="image name:tag (e.g. centos:7.2)", required=True)
    runtime_replace_parser.add_argument("-l", "--language", dest="language",
                                    help="language type, currently support python, python3 and r. " +
                                    "Used to auto-configure container command.", required=True, choices=['r','python', 'python3'])
    runtime_replace_parser.add_argument("-v", "--volume", dest="shared", action="append",
                                    help="Bind mount a volume, can work with multiple volumes" +
                                    " format: HOST:CONTAINER:[rw|ro]")
    runtime_replace_parser.add_argument("-s", "--setting", dest="setting", action="append",
                                    help="Setting for the runtime. Allow call multiple times." +
                                    " (e.g. memory_mb=256). See more in document.")
    # For runtime-show
    runtime_show_parser.add_argument("-r", "--runtime", dest="runtime_id",
                                    help="show specific runtime with given runtime id in configuration file")

    # For runtime-delete
    runtime_delete_parser.add_argument("-r", "--runtime", dest="runtime_id",
                                    help="unique string which presents configurations for a container runtime)", required=True)

    # For runtime-edit
    runtime_edit_parser.add_argument("-e", "--editor", dest="editor",
                                    help="editor command (e.g. vim, etc)", required=False)

    # For runtime-backup
    runtime_backup_parser.add_argument("-f", "--file", dest="file",
                                    help="The file that the runtime configurations will be saved (xml format)", required=True)

    # For runtime-restore
    runtime_restore_parser.add_argument("-f", "--file", dest="file",
                                    help="The file that includes the runtime configurations", required=True)


    options = parser.parse_args()

    return options

def distribute_docker_file(pool, d_file, dest, hosts):
    if dest is None:
        dest = os.path.join(TMP, d_file)
    for host in hosts:
        cmd = Rsync('Distribute docker image file.',
                    d_file,
                    dest,
                    None,
                    host.strip())
        pool.addCommand(cmd)
    pool.join()
    pool.check_results()

def remote_cmd(pool, cmd_str, hosts):
    for host in hosts:
        cmd = Command('Execute remote cmd on all hosts.',
                        cmd_str,
                        ctxt = REMOTE,
                        remoteHost = host.strip())
        pool.addCommand(cmd)
    pool.join()

    items = pool.getCompletedItems()
    results = None
    for i in items:
        if i.results.rc or i.results.halt or not i.results.completed:
            logger.error("error running command: %s" % cmd_str)
            raise Exception("error running command")
        if results is None:
            results = items[0].get_results().stdout
    return results

def get_docker_image(url):
    '''
    download docker image from given url
    '''
    logger.info("Getting docker image from %s" % url)
    (ok, out) = subprocess.getstatusoutput('docker pull %s' % url)
    if ok != 0:
        logger.error("Cannot pull docker image from %s" % url)
        sys.exit(1)
    name_tag = url.split("/")[-1]
    image_name = "".join([name_tag.replace(":", "_"), TAR_GZ])
    image_path = os.path.join(TMP, image_name)
    logger.info("Saving docker image to %s" % image_path)
    (ok, out) = subprocess.getstatusoutput('docker save -o %s %s' % (image_path, url))
    if ok != 0:
        logger.error("Cannot save docker image %s to %s" % (url, image_path))
        sys.exit(1)

    return image_path

def install_image_remote(hosts, pool, docker_image):
    try:
        step = 0
        logger.info('Checking whether docker is installed on all hosts...')
        check_docker_cmd = 'which docker'
        remote_cmd(pool, check_docker_cmd, hosts)
        step = 1

        logger.info('Distributing image file %s to all hosts...' % docker_image)
        dest = os.path.join(DBHOME, TMP, os.path.basename(docker_image))
        distribute_docker_file(pool, docker_image, dest, hosts)
        step = 2

        logger.info('Loading image on all hosts...')
        docker_load_cmd = " ".join(["docker load -i ", dest])
        remote_cmd(pool, docker_load_cmd, hosts)
        step = 3

        logger.info('Removing temporary image files on all hosts...')
        docker_load_cmd = " ".join(["rm -f ", dest])
        remote_cmd(pool, docker_load_cmd, hosts)
    except Exception as ex:
        if (step == 0):
            logger.error("Cannot find docker in all hosts, please check docker installation again")
        elif (step == 1):
            logger.error("Unable to distribute the docker tarball into all hosts")
        elif (step == 2):
            logger.error("Cannot load docker image in all hosts")
        elif (step == 3):
            logger.error("Unable to remove temporary files")
        raise(ex)

def install_image_local(hosts, pool, docker_file):
    try:
        step = 0
        logger.info('Checking whether docker is installed on all hosts...')
        check_docker_cmd = 'which docker'
        remote_cmd(pool, check_docker_cmd, hosts)
        step = 1

        logger.info('Checking whether docker image tarball is exsited on all hosts...')
        check_file_cmd = " ".join(["ls ", docker_file])
        remote_cmd(pool, check_file_cmd, hosts)
        step = 2

        logger.info('Loading image on all hosts via local file...')
        docker_load_cmd = " ".join(["docker load -i ", docker_file])
        remote_cmd(pool, docker_load_cmd, hosts)
    except Exception as ex:
        if (step == 0):
            logger.error("Cannot find docker in all hosts, please check docker installation again")
        elif (step == 1):
            logger.error("Docker image tarball %s is not existed on all hosts, please copy them again" % docker_file)
        elif (step == 2):
            logger.error("Cannot load docker image in all hosts")
        raise(ex)

def install_images(options, pool, hosts):
    '''
    install docker image either form url or file
    '''
    if options.local:
        if options.file:
            docker_image = options.file
        else:
            logger.error("Please use -f option to indicate file location")
            sys.exit(1)
    if options.file:
        docker_image = options.file
    if options.url:
        docker_image = get_docker_image(options.url)

    if options.local:
        install_image_local(hosts, pool, docker_image)
    else:
        if not os.path.exists(docker_image):
            logger.error("docker image file %s does not exist" % docker_image)
            sys.exit(1)
        install_image_remote(hosts, pool, docker_image)

def delete_images(options, pool, hosts):
    '''
    delete docker image
    '''
    docker_image = options.image_name

    try:
        logger.info('Checking whether docker is installed on all hosts...')
        check_docker_cmd = 'which docker'
        remote_cmd(pool, check_docker_cmd, hosts)

        logger.info('Deleting image %s on all hosts...' % docker_image)
        docker_delete_cmd = " ".join(["docker rmi -f ", docker_image])
        remote_cmd(pool, docker_delete_cmd, hosts)
    except Exception as ex:
        raise(ex)

def update_xml_check_shared_folders(runtime, elements):
    '''
    update shared folder settings
    '''
    for shareDir in runtime.findall('shared_directory'):
        runtime.remove(shareDir)

    dest = os.path.join(DBHOME, PLCONTAINER_CLIENT_DIR)

    shareDir = ElementTree.Element("shared_directory",
            host=dest, container="/clientdir",
            access="ro")
    runtime.append(shareDir)

    if 'hasShared' in elements:
        count = 0
        for plHost in elements['host']:
            shareDir = ElementTree.Element("shared_directory",
                    host=plHost, container=elements['container'][count],
                    access=elements['access'][count])
            count = count + 1
            runtime.append(shareDir)


def update_xml_check_settings(runtime, elements):
    '''
    update other settings.
    '''
    if not 'hasSetting' in elements:
        return

    for setting in runtime.findall('setting'):
        runtime.remove(setting)

    for setting in elements['setting']:
        xml_setting = ElementTree.Element("setting")
        xml_setting.set(setting, elements['setting'][setting])
        xml_setting.text = ""
        runtime.append(xml_setting)

def update_xml_runtime(runtime, elements):
    '''
    update a runtime configuration.
    '''
    image = runtime.find('image')
    image.text = elements['image']
    command = runtime.find('command')
    if elements['language'] == "r":
        command.text = "/clientdir/rclient.sh"
    elif elements['language'] == "python":
        command.text = "/clientdir/pyclient.sh"
    else:
        command.text = "/clientdir/py3client.sh"

    update_xml_check_shared_folders(runtime, elements)
    update_xml_check_settings(runtime, elements)


def update_xml_new_runtime(elements):
    '''
    Create a new runtime configuration.
    '''
    runtime = ElementTree.Element("runtime")
    runtime_id = ElementTree.Element("id")
    runtime_id.text = elements['id']
    image = ElementTree.Element("image")
    image.text = elements['image']

    # default command
    command = ElementTree.Element("command")
    if elements['language'] == "r":
        command.text = "/clientdir/rclient.sh"
    elif elements['language'] == "python":
        command.text = "/clientdir/pyclient.sh"
    else:
        command.text = "/clientdir/py3client.sh"
    runtime.append(runtime_id)
    runtime.append(image)
    runtime.append(command)

    update_xml_check_shared_folders(runtime, elements)
    update_xml_check_settings(runtime, elements)

    return runtime

def update_reformat_xml(xmltree):
    '''
    reformat xml to a pretty print format
    '''
    xmlString = ElementTree.tostring(xmltree)
    reformatString = xml.dom.minidom.parseString(xmlString)
    return reformatString.toprettyxml(indent = '    ') 

def add_xml_config(filename, elements):
    '''
    Update the xml file with the new runtime configuration. Fail if the given tag exists.
    '''
    try:
        xmltree = ElementTree.parse(filename)
        logger.debug("file name is %s" % (filename))
        root = xmltree.getroot()

        for runtime in root.findall('runtime'):
            runtime_id = runtime.find('id').text
            if runtime_id == elements['id']:
                logger.error("runtime configuration exists for %s. You might want the runtime-replace command?" % runtime_id)
                raise Exception("runtime configuration exists")

        runtime = update_xml_new_runtime(elements)
        root.append(runtime)

        newXMLTreeString = update_reformat_xml(root)
        newXMLTreeString = os.linesep.join([s for s in newXMLTreeString.splitlines() if not s.strip() == ''])
        xmlfile = open(filename, "w")
        xmlfile.write(newXMLTreeString)
        xmlfile.close()
    except Exception as ex:
        logger.error("Update XML file Error")
        logger.error(str(ex))
        raise Exception("Update XML file Error")

def update_xml_config(filename, elements):
    '''
    Update the xml file by given elements dict, if the given tag is
    not existed, then create a new tag. Otherwise, update the current
    tag.
    '''
    isOld = False
    try:
        xmltree = ElementTree.parse(filename)
        logger.debug("file name is %s" % (filename))
        root = xmltree.getroot()

        for runtime in root.findall('runtime'):
            runtime_id = runtime.find('id').text
            if runtime_id == elements['id']:
                isOld = True
                break

        if not isOld:
            logger.info("runtime configuration does not exist. Adding a new one.")
            add_xml_config(filename, elements)
            return

        update_xml_runtime(runtime, elements)

        newXMLTreeString = update_reformat_xml(root)
        newXMLTreeString = os.linesep.join([s for s in newXMLTreeString.splitlines() if not s.strip() == ''])
        xmlfile = open(filename, "w")
        xmlfile.write(newXMLTreeString)
        xmlfile.close()
    except Exception as ex:
        logger.error("Update XML file Error")
        logger.error(str(ex))
        raise Exception("Update XML file Error")

def delete_xml_config(filename, del_runtime_id):
    '''
    Delete given runtime from the xml file.
    '''
    try:
        xmltree = ElementTree.parse(filename)
        logger.debug("file name is %s" % (filename))
        root = xmltree.getroot()

        for runtime in root.findall('runtime'):
            runtime_id = runtime.find('id').text
            if runtime_id == del_runtime_id:
                root.remove(runtime)
                newXMLTreeString = update_reformat_xml(root)
                newXMLTreeString = os.linesep.join([s for s in newXMLTreeString.splitlines() if not s.strip() == ''])
                xmlfile = open(filename, "w")
                xmlfile.write(newXMLTreeString)
                xmlfile.close()
                return True
    except Exception as ex:
        logger.error("Update XML file Error")
        logger.error(str(ex))
        raise Exception("Update XML file Error")
    return False

def get_file_md5(file):
    '''
    Generates an MD5 hash of a file.  The hash returned is an array
    of bytes and cannot be printed directly to the screen or log file
    without converting to hex.
    '''
    with open(file, 'rb') as f:
        return hashlib.md5(f.read()).digest()


def copy_file_to_tmp(host, file):
    """
    Copies file to a local temp file and returns the temp file name.
    While mktemp() is not "safe", the possibilities of this utility
    being run multiple times on the same file and getting the same
    temp filename is so remote that this is acceptable.
    """
    tmp_file = tempfile.mktemp()
    if not os.path.exists(file):
        default_config = os.path.join(DBHOME, PLCONTAINER_GPDB_DIR, PLCONTAINER_CONFIG_FILE)
        cmd = Rsync('copy conf file from GPDB folder', default_config, file, host)
        cmd.run()
    cmd = Rsync('copy conf file', file, tmp_file, host)
    cmd.run()
    if not os.path.exists(tmp_file):
        logger.error("Error while copying configuration file. Check that PL/Container is initialized" +
                     " or use --reset to initialize with default configuration")
        tmp_file = None
    return tmp_file


def validate_xml(filename):
    image_list=[]
    for line in os.popen("docker images | awk 'NR>1 {print $1 \":\" $2}'"):
        line = line.strip('\n')
        image_list.append(line)
    correct = True
    try:
        xmltree = ElementTree.parse(filename)
        root = xmltree.getroot()
        for runtime in root.findall('runtime'):

            # Only one 'id' subelement is allowed.
            pattern = re.compile(r'^[a-zA-Z0-9][a-zA-Z0-9_.-]*$')
            num = len(runtime.findall('id'))
            if num < 1:
                logger.error("Lack of 'id' subelement in a runtime element.")
                raise Exception("Validation failed")
            if num > 1:
                logger.error("Found %d (1 is legal) 'id' subelements in a runtime element. The first one 'id': %s", num, runtime.find('id').text)
                raise Exception("Validation failed")

            runtime_id = runtime.find('id').text

            if len(runtime_id) >= RUNTIME_ID_MAX_LENGTH:
                logger.error("length of runtime id %s should not exceeds 63.", runtime_id)
                raise Exception("Validation failed")
            if  not pattern.match(runtime_id):
                logger.error("'id' name with wrong format: %s", runtime_id)
                raise Exception("Validation failed")

            # Only one 'image' subelement is allowed.
            num = len(runtime.findall('image'))
            if num < 1:
                logger.error("Lack of 'image' subelement in a runtime element in runtime %s.", runtime_id)
                raise Exception("Validation failed")
            if num > 1:
                logger.error("Found %d (1 is legal) 'image' subelements in a runtime %s. The first one: %s", num, runtime_id, runtime.find('image').text)
                raise Exception("Validation failed")
            # FIXME: Need to check the naming convention?

            if not runtime.find('image').text in image_list:
                logger.warning("image %s is not in the list of 'docker images'" % runtime.find('image').text)

            # Only one 'command' subelement is allowed.
            num = len(runtime.findall('command'))
            if num < 1:
                logger.error("Lack of 'command' subelement in a runtime element in runtime %s.", runtime_id)
                raise Exception("Validation failed")
            if num > 1:
                logger.error("Found %d (1 is legal) 'command' subelements in runtime %s. The first one 'command': %s", num, runtime_id, runtime.find('command').text)
                raise Exception("Validation failed")

            # Shared
            container_pathlist = []
            for share_folder in runtime.findall('shared_directory'):

                if 'access' in share_folder.attrib and share_folder.attrib['access'] != 'ro' and share_folder.attrib['access'] != 'rw':
                    logger.error("access should be 'ro' or 'rw' in shared_directory in runtime %s, but now: %s", runtime_id, share_folder.attrib['access'])
                    raise Exception("Validation failed")

                if  sorted(list(share_folder.attrib.keys())) != ['access', 'container', 'host']:
                    logger.error("Unknown or missing attributes (need and only need 'access', 'host' and 'container') in shared_directory in runtime %s", runtime_id)
                    logger.error("Current: %s", list(share_folder.attrib.keys()))
                    raise Exception("Validation failed")

                # This path is reserved for unix domain socket connection.
                if share_folder.attrib['container'] == '/tmp/plcontainer' or share_folder.attrib['container'] == '/tmp/plcontainer/':
                    logger.error("container path should not be /tmp/plcontainer (reserved for unix domain socket connection) in runtime %s", runtime_id)
                    raise Exception("Validation failed")

                container_pathlist.append(share_folder.attrib['container'])

            tmp_list = list(set(container_pathlist))
            if len(tmp_list) != len(container_pathlist):
                logger.error("container should not have same mount points (check shared_directory) in runtime %s", runtime_id)
                logger.error("Note if you are modifying/adding runtime configuration please do not use the reserved container paths /clientdir and /tmp/plcontainer")
                raise Exception("Validation failed")

            # Settings
            for settings in runtime.findall('setting'):
                if 'memory_mb' in settings.attrib:
                    memory_mb_str = settings.attrib['memory_mb']
                    try:
                        memory_mb = int(memory_mb_str)
                        if memory_mb <= 0:
                            logger.error("memory_mb should > 0 in runtime %s, but now: '%s'", runtime_id, memory_mb_str)
                            raise Exception("Validation failed")
                    except ValueError:
                        logger.error("memory_mb should be a positive integer in runtime %s, but now: '%s'", runtime_id, memory_mb_str)
                        raise Exception("Validation failed")
                elif 'cpu_share' in settings.attrib:
                    cpu_share_str = settings.attrib['cpu_share']
                    try:
                        cpu_share = int(cpu_share_str)
                        if cpu_share <= 0:
                            logger.error("cpu_share should > 0 in runtime %s, but now: '%s'", runtime_id, cpu_share_str)
                            raise Exception("Validation failed")
                    except ValueError:
                        logger.error("cpu_share should be a positive integer in runtime %s, but now: '%s'", runtime_id, cpu_share_str)
                        raise Exception("Validation failed")
                elif 'use_container_logging' in settings.attrib:
                    use_container_logging_str = settings.attrib['use_container_logging'].lower()
                    if use_container_logging_str != 'yes' and use_container_logging_str != 'no':
                        logger.error("'use_container_logging' should be 'yes' or 'no' in runtime %s, but now: '%s'", runtime_id, use_container_logging_str)
                        raise Exception("Validation failed")
                elif 'enable_network' in settings.attrib:
                    enable_network_str = settings.attrib['enable_network'].lower()
                    if enable_network_str != 'yes' and enable_network_str != 'no':
                        logger.error("'enable_network' should be 'yes' or 'no' in runtime %s, but now: '%s'", runtime_id, enable_network_str)
                        raise Exception("Validation failed")
                elif 'resource_group_id' in settings.attrib:
                    resource_group_id_str = settings.attrib['resource_group_id']
                    if resource_group_id_str.isdigit() != True:
                        logger.error("'resource_group_id_str' cannot contain non-digit, it should be oid of resource group")
                        raise Exception("Validation failed")
                elif 'roles' in settings.attrib:
                    roles_str = settings.attrib['roles']
                    if not roles_str:
                        logger.error("'roles_str' is empty, please add roles (separate with ',') or remove this element")
                        raise Exception("Validation failed")
                else:
                    logger.error("Found illegal setting in runtime %s: %s", runtime_id, settings.attrib)
    except Exception as ex:
        correct = False 
        logger.error("XML file does not conform XML specification or wrong settings.")
        logger.error(str(ex))
    return correct

def edit_file(filename, editor):
    """
    Starts up options.editor (default is vi) to edit the file.
    """
    edit = True
    while edit:
        distribute = True
        h1 = get_file_md5(filename)
        os.system('%s %s' % (editor, filename))
        h2 = get_file_md5(filename)
        if h1 == h2:
            logger.info("File was not changed, no configuration update is done")
            distribute = False
        validxml = validate_xml(filename)
        if validxml:
            edit = False
        else:
            # If the file is invalid we won't distribute it
            distribute = False
            edit = userinput.ask_yesno(None, "Continue editing the file?", 'N')
    return distribute


def distribute_file(pool, src_filename, dest_filename, gpdb_locations):
    """
    Distributes the edited file to all GPDB instances
    """
    for l in gpdb_locations:
        dest_file = os.path.join(l[1], dest_filename)
        logger.debug("Distributing to %s:%s", l[0], dest_file)
        cmd = Rsync('distribute changed file',
                  src_filename,
                  dest_file,
                  None,
                  l[0])
        pool.addCommand(cmd)
    pool.join()
    pool.check_results()


def backup_file(pool, gpdb_locations, file):
    """
    Backs up a file to .bak
    """
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    dest_filename = '.' + file + '.bak' + timestamp
    logger.info("Old configuration is saved as '%s'" % dest_filename)
    for l in gpdb_locations:
        src_file = os.path.join(l[1], file)
        dest_file = os.path.join(l[1], dest_filename)
        cmd = Rsync('backup file',
                  src_file,
                  dest_file,
                  l[0],
                  l[0])
        pool.addCommand(cmd)
    pool.join()
    pool.check_results()

def get_all_hosts(array):
    hosts = set()
    for hostname in array.get_hostlist():
       hosts.add(hostname)
    return hosts

def get_segment_locations(seg_list):
    """
    Returns a list of tuples that are of the format (hostname, gpdb directory)
    for every segment instance in the cluster.
    """
    gpdb_locations = []
    for s in seg_list:
        gpdb_locations.append((s.getSegmentHostName(), s.getSegmentDataDirectory()))
    return gpdb_locations

def remote_remove_files(pool, gpdb_locations, remote_filename):
    """
    Removes the specified file from all the hosts of GPDB cluster
    """
    for l in gpdb_locations:
        remote_file = os.path.join(l[1], remote_filename)
        cmd = Command('Remove %s:%s' % (l[0], remote_file),
                      'rm -f %s || exit 0' % remote_file,
                      ctxt = REMOTE,
                      remoteHost = l[0])
        pool.addCommand(cmd)
    pool.join()
    pool.check_results()


def show_file(filename):
    f = open(filename, 'rb')
    sys.stdout.write('%s\n' % f.read())
    f.close()

def runtime_show_xml(filename, r_id):
    '''
    Print configurations
    :param filename:  
    :return: 
    '''
    try:
        no_entry = True
        xmltree = ElementTree.parse(filename)
        root = xmltree.getroot()
        sys.stdout.write("PL/Container Runtime Configuration: \n")
        for runtime in root.findall('runtime'):
            runtime_id = runtime.find('id').text
            if (r_id is None or r_id == runtime_id):
                no_entry = False
                sys.stdout.write("---------------------------------------------------------\n")
                sys.stdout.write("  Runtime ID: %s\n" % runtime_id)
                image_name = runtime.find('image').text
                sys.stdout.write("  Linked Docker Image: %s\n" % image_name)
                sys.stdout.write("  Runtime Setting(s): \n")
                for settings in runtime.findall('setting'):
                    if 'memory_mb' in settings.attrib:
                        sys.stdout.write("  ---- Container Memory Limited: %s MB\n" % settings.attrib['memory_mb'])
                    elif 'cpu_share' in settings.attrib:
                        sys.stdout.write("  ---- Container CPU share: %s\n" % settings.attrib['cpu_share'])
                    elif 'use_container_logging' in settings.attrib:
                        sys.stdout.write("  ---- Use Container Logging: %s\n" % settings.attrib['use_container_logging'])
                    elif 'resource_group_id' in settings.attrib:
                        sys.stdout.write("  ---- Resource Group ID: %s\n" % settings.attrib['resource_group_id'])
                    elif 'roles' in settings.attrib:
                        sys.stdout.write("  ---- Allowed Roles: %s\n" % settings.attrib['roles'])
                sys.stdout.write("  Shared Directory: \n")
                for share_folder in runtime.findall('shared_directory'):
                    sys.stdout.write("  ---- Shared Directory From HOST '%s' to Container '%s', access mode is '%s'\n" % (share_folder.attrib['host'], share_folder.attrib['container'], share_folder.attrib['access']))
                sys.stdout.write("---------------------------------------------------------\n")
                sys.stdout.write("\n")

        if no_entry:
            sys.stdout.write("No runtime or no specified runtime was found.\n")
    except Exception as ex:
        logger.error("XML file does not conform XML specification or no related element was found.")
        logger.error(str(ex))
        raise Exception("runtime-show error")

def get_conf_settings(options, elements):
    '''
    get the settings for the runtime configuration.
    '''
    elements['hasSetting'] = True
    elements['setting'] = dict()
    for setting in options.setting:
        strList = setting.split("=")
        if len(strList) != 2:
            raise Exception("Bad setting format: %s" % setting)
        if strList[0] not in ("memory_mb", "cpu_share", "use_container_logging", "resource_group_id", "roles", "enable_network"):
            raise Exception("Bad setting key: %s" % strList[0])
        elements['setting'][strList[0]] = strList[1]

def get_conf_shared_dirs(options, elements):
    '''
    get the shared dirs
    '''
    elements['hasShared'] = True
    shareStrs = options.shared
    elements['host'] = list()
    elements['container'] = list()
    elements['access'] = list()
    for shareStr in shareStrs:
        strList = shareStr.split(":")
        elements['host'].append(strList[0])
        elements['container'].append(strList[1])
        if len(strList) >= 3:
            if strList[2] == "ro" or strList[2] == "rw":
                elements['access'].append(strList[2])
            else:
                raise Exception("Access type can only be \"ro\" or \"rw\". ")
        else:
            elements['access'].append("ro")

def get_conf_elements(options):
    '''
    build an element dict for configuration
    '''
    elements = dict()
    elements['id'] = options.runtime_id
    elements['image'] = options.image_name
    elements['language']=options.language.lower()
    if options.shared:
       get_conf_shared_dirs(options, elements)
    if options.setting:
       get_conf_settings(options, elements)
    return elements

def checkAllFileMD5(pool, gpdb_locations):
    '''
    
    :param pool: workpool 
    :param gpdb_locations: gpdb data dir
    :return TRUE if match, FALSE if not match: 
    '''
    logger.info("Verifying...")
    for l in gpdb_locations:
        remote_file = os.path.join(l[1], PLCONTAINER_CONFIG_FILE)
        cmd = Command('Check %s:%s' % (l[0], remote_file),
                  'md5sum %s || exit 0' % remote_file,
                  ctxt=REMOTE,
                  remoteHost=l[0])
        pool.addCommand(cmd)
    pool.join()

    items = pool.getCompletedItems()
    results = None
    for i in items:
        if i.results.rc or i.results.halt or not i.results.completed:
            logger.error("error running command")
            raise Exception("error running command")
        # Put the QD conf file MD5 into resluts and do comparison
        if results is None:
            results = i.get_results().stdout.split(" ")[0]
        else:
            if results != i.get_results().stdout.split(" ")[0]:
                return False
    return True


def main():
    retval = 0
    pool = None
    tmp_file = None
    options = None
    new_config = False

    global DBHOME
    global DBTYPE

    options = parseargs()
    if options.verbose:
        gplog.enable_verbose_logging()

    try:
        DBHOME = os.environ.get('GPHOME')
        if not DBHOME:
            DBHOME = os.environ.get('PGHOME')
            if not DBHOME:
                raise Exception('$GPHOME or $PGHOME environment variable is not set')
            else:
                DBTYPE = "pg"
        else:
            DBTYPE = "gp"

        logger.debug('Retrieving GPDB configuration from database...')
        url = dbconn.DbURL()
        
        #for pg
        hosts = set()
        gpdb_locations = []

        if DBTYPE == "gp":
            array = gparray.GpArray.initFromCatalog(url, True)
            segs = array.getDbList()
            hosts = get_all_hosts(array)
            gpdb_locations = get_segment_locations(segs)
        else:
            hosts.add('localhost')   #todo 
            pgdata_path = os.environ.get('PGDATA')
            if not pgdata_path:
                raise Exception('$PGDATA environment variable is not set for postgresql.')
            gpdb_locations.append( ('localhost', pgdata_path) )

            
        logger.debug('GPDB Configuration is:')

        # FIXME: do them for configure commands only.
        remote_dir = gpdb_locations[0][1]
        remote_host = gpdb_locations[0][0]
        remote_filename = os.path.join(remote_dir, PLCONTAINER_CONFIG_FILE)
        tmp_file = copy_file_to_tmp(remote_host, remote_filename)

        for host, directory in gpdb_locations:
            logger.debug('    host "%s" directory "%s"' % (host, directory))

        if options.subparser_name == "image-add":
            logger.debug('Preparing the worker pool...')
            pool = WorkerPool(PL_BATCH_SIZE)
            install_images(options, pool, hosts)

        if options.subparser_name == "image-delete":
            logger.debug('Preparing the worker pool...')
            pool = WorkerPool(PL_BATCH_SIZE)
            delete_images(options, pool, hosts)

        if options.subparser_name == "image-list":
            os.system("docker images")

        if options.subparser_name == "runtime-add":
            elements = get_conf_elements(options)
            add_xml_config(tmp_file, elements)
            valid_xml = validate_xml(tmp_file)
            if not valid_xml:
                retval = 1
            else:
                logger.debug('Preparing the worker pool...')
                pool = WorkerPool(PL_BATCH_SIZE)
                logger.info('Distributing file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                distribute_file(pool, tmp_file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                new_config = True

        if options.subparser_name == "runtime-replace":
            elements = get_conf_elements(options)
            update_xml_config(tmp_file, elements)
            valid_xml = validate_xml(tmp_file)
            if not valid_xml:
                retval = 1
            else:
                logger.debug('Preparing the worker pool...')
                pool = WorkerPool(PL_BATCH_SIZE)
                logger.info('Distributing file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                distribute_file(pool, tmp_file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                new_config = True

        if options.subparser_name == "runtime-show":
            if options.runtime_id is None:
                runtime_show_xml(tmp_file, None)
            else:
                runtime_show_xml(tmp_file, options.runtime_id)

        if options.subparser_name == "runtime-delete":
            if delete_xml_config(tmp_file, options.runtime_id):
                logger.debug('Preparing the worker pool...')
                pool = WorkerPool(PL_BATCH_SIZE)
                logger.info('Distributing file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                distribute_file(pool, tmp_file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                new_config = True
            else:
                logger.info('This runtime configuraiton does not exist. Ignoring.')

        if options.subparser_name == "runtime-edit":
            if options.editor is None:
                options.editor  = os.getenv("EDITOR", "vi") # default editer is 'vi'
            changed = edit_file(tmp_file, options.editor)
            if changed:
                logger.debug('Preparing the worker pool...')
                pool = WorkerPool(PL_BATCH_SIZE)
                logger.info('Distributing file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                distribute_file(pool, tmp_file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                new_config = True

        if options.subparser_name == "runtime-backup":
            fb = open(tmp_file, 'rb')
            fw = open(options.file, 'wb')
            fw.write(fb.read())
            fw.write(b'\n')
            fw.close()
            fb.close()

        if options.subparser_name == "runtime-restore":
            valid_xml = validate_xml(options.file)
            if not valid_xml:
                logger.error('Bad xml format. Please fix the file at first.')
                retval = 1
            else:
                pool = WorkerPool(PL_BATCH_SIZE)
                logger.info('Distributing file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                distribute_file(pool, options.file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                new_config = True

        if options.subparser_name == "runtime-verify":
            logger.debug('Preparing the worker pool...')
            valid_xml = validate_xml(tmp_file)
            if not valid_xml:
                logger.error('Bad xml format. Please fix them using the runtime-edit sub-command.')
                retval = 1
            else:
                pool = WorkerPool(PL_BATCH_SIZE)
                verified = checkAllFileMD5(pool, gpdb_locations)
                if not verified:
                    logger.error('Detected configuration files are not the same in all hosts!')
                    edit = userinput.ask_yesno(None, "Fix it by copying master configuration file?", 'N')
                    if edit:
                        logger.info('Distributing master configuration file %s to all locations...' % PLCONTAINER_CONFIG_FILE)
                        distribute_file(pool, tmp_file, PLCONTAINER_CONFIG_FILE, gpdb_locations)
                        new_config = True
                    else:
                        logger.error('Verify Failed!')
                        retval = 1
                else:
                    logger.info('Verify success!')

        if new_config:
            logger.info('Configuration has changed. Run "select * from plcontainer_refresh_config" in open sessions.' +
                        ' New sessions will get new configuration automatically.')

    except Exception as ex:
        retval = 1
        logger.fatal("%s failed. (Reason='%s') exiting..." % (execname, ex))
        if options.verbose:
            logger.exception(ex)

    finally:
        if tmp_file:
            os.unlink(tmp_file)
        if pool:
            logger.debug('Shutting down workers...')
            pool.haltWork()
            pool.joinWorkers()
        return retval

if __name__ == '__main__':
    execname = os.path.split(__file__)[-1]
    gplog.setup_tool_logging(execname, getLocalHostname(), getUserName())
    logger = gplog.get_default_logger()
    sys.exit(main())
