#
# Copyright (C) 2016 Rolf Neugebauer <rolf.neugebauer@docker.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#

"""Main entry point to run the regression test on a remote host"""

import abc
import argparse
import base64
import fnmatch
import os
import shutil
import socket
import stat
import subprocess
import sys
import tarfile
import tempfile
import threading

import rt.httpsvr
import rt.log
import winrm

DEVNULL = open(os.devnull, 'wb')

# Temporary directories on the remote hosts
_WIN_TMP_DIR = "/c/Users/%s/AppData/Local/RT"
_OSX_TMP_DIR = "/tmp/RT"

# globals for ports to use as well as IP address to pass to remotes
_PORT_AUX_HTTP = -1
_PORT_AUX_LOG = -1
_LOCAL_IP = None
_KEYFILE = None


def _targz_dir(dst, directory, excludes=None):
    """tar up a directory"""
    local_root, dirname = os.path.split(directory)
    tf = tarfile.open(dst, mode="w:gz")
    for root, _, files in os.walk(directory):
        for filename in files:
            fp = os.path.join(root, filename)
            if excludes is not None:
                if any(fnmatch.fnmatch(fp, pat) for pat in excludes):
                    continue
            archive_root = os.path.relpath(root, local_root)
            archive_path = os.path.join(archive_root, filename)
            tf.add(fp, arcname=archive_path)
    tf.close()


class Unbuffered(object):
    """
    This class is used to make sure stdout is not buffered
    """
    def __init__(self, stream):
        self.stream = stream

    def write(self, data):
        self.stream.write(data)
        self.stream.flush()

    def __getattr__(self, attr):
        return getattr(self.stream, attr)


class _Host:
    """An abstract base class to interact with a host"""
    __metaclass__ = abc.ABCMeta

    def __init__(self, remote, user, password):
        """Initialise the class"""
        self.remote = remote
        self.user = user
        self.password = password
        return

    @abc.abstractmethod
    def run_sh(self, cmd):
        """Run a shell command on Host"""
        pass

    def cp_to(self, lpath, rpath, excludes=None):
        """Copy a file from the local host to the remote host. If @lpath
        is a directory then copy recursively. @rpath *must* be a
        directory. Optionally supply a list of glob expressions for
        files to ignore."""
        print("COPY TO: %s -> %s" % (lpath, rpath))
        if not os.path.isdir(lpath):
            self._cp_to(lpath, rpath)
        # we have a directory to copy.
        # 1. Tar up directory to a temporary file
        fd, fp = tempfile.mkstemp(suffix=".tar.gz")
        _targz_dir(fp, lpath, excludes)
        _, fn = os.path.split(fp)

        # 2. Copy file to the host and remove tmp file
        rc = self._cp_to(fp, rpath)
        os.close(fd)
        os.remove(fp)
        if not rc == 0:
            return rc

        # 3. Untar the file on the host
        rc = self.run_sh("(cd %s; tar xzf %s; rm %s)" %
                         (rpath, fn, fn))
        return rc

    @abc.abstractmethod
    def _cp_to(self, lpath, rpath):
        """ Class specfic implementation of a copy to operation"""
        pass

    @abc.abstractmethod
    def cp_from(self, rpath, lpath):
        """Copy a directory from the remote host to the local host. If @rpath
        ends with a '/' copy the directory recursively. @lpath *must*
        be a directory."""
        pass

    def run_rt_local(self, rpath, rbin, args):
        """Run tests on remote system. This is special method as it allows
        different backends to modify the arguments a little if
        needed. @rpath is where the test cases can be found.
        @rbin is the name of the rt-local binary"""

        rc = self.run_sh("cd %s; %s %s" % (rpath, rbin, args))
        return rc


class _LogThread(threading.Thread):
    """Listen on a logging socket and print stuff till close()"""

    def __init__(self, port):
        threading.Thread.__init__(self)
        self.port = port

    def run(self):
        rt.log.log_svr(self.port)


class WinRMHost(_Host):
    """A class to interact with remote host via Windows Remote Management"""
    # This is fun.
    # - WinRM does not allow one to send file to and fro. So we use
    #   WinRM to execute bash commands on host and use a local web
    #   server to handle GET/POST request to transfer file.  When
    #   transferring directories we tar them before transferring.
    # - WinRM also only gives you the output of a command once it is
    #   done.  So, we log over a socket and have a log server running
    #   locally to print the logs as they come in.
    # - pywinrm chokes when a command runs to long, so never returns.
    #   So we start in the background and exit once the log thread exits.

    def __init__(self, remote, user, password):
        """Initialise the class"""
        super(WinRMHost, self).__init__(remote, user, password)

        self.our_ip = _LOCAL_IP if _LOCAL_IP else self.get_local_ip(remote)

        self.session = winrm.Session("%s" % remote,
                                     auth=(self.user, self.password),
                                     server_cert_validation='ignore')

    @staticmethod
    def get_local_ip(remote):
        """Try to get local IP address used to connect to the remote host"""
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect((remote, 53))
            our_ip = s.getsockname()[0]
            s.close()
        except:
            our_ip = socket.gethostbyname(socket.gethostname())
        return our_ip

    def run_sh(self, cmd):
        """Run a shell command on the host"""

        # This is really, really tedious, but I could not get cmd.exe
        # to run bash directly, so we start powershell and run it from
        # there. Further, winrm.run_ps does not work either.
        s = self.session
        print("EXEC: %s" % cmd)
        rs = s.run_cmd("powershell.exe",
                       ['-NonInteractive',
                        '-ExecutionPolicy', 'Unrestricted',
                        '-Command',
                        "\"& 'C:\\Program Files\\Git\\bin\\bash.exe' " +
                        "--login -c '%s' \"" % cmd])
        if len(rs.std_out.strip()):
            print("STDOUT:")
            for l in rs.std_out.split('\n'):
                print(l)
        if len(rs.std_err.strip()):
            print("STDERR:")
            for l in rs.std_err.split('\n'):
                print(l)
        return rs.status_code

    def _cp_to(self, lpath, rpath):
        """Copy a file in @lpath to the @rpath directory on the host."""
        _, fn = os.path.split(lpath)
        # 1. Start web server with file
        s = rt.httpsvr.dl_svr_start(_PORT_AUX_HTTP, lpath)
        # 2. Execute curl (after making sure the dir exist)
        self.run_sh("mkdir -p %s" % rpath)
        rc = self.run_sh("(cd %s; curl -s %s:%d -o %s)" %
                         (rpath, self.our_ip, _PORT_AUX_HTTP, fn))
        # 3. Make sure server is stopped
        rt.httpsvr.dl_svr_wait(s)
        return rc

    def _cp_from(self, rpath, lpath):
        """Copy a file from @rpath to the local @lpath directory."""
        rc = subprocess.call("mkdir -p %s" % lpath, shell=True)
        fp, fn = os.path.split(rpath)
        # 1. Start web server to write a POST to the file
        s = rt.httpsvr.ul_svr_start(_PORT_AUX_HTTP, os.path.join(lpath, fn))
        # 2. Execute curl
        rc = self.run_sh("(cd %s; curl -s -X POST --data-binary @%s %s:%d)" %
                         (fp, fn, self.our_ip, _PORT_AUX_HTTP))
        # 3. Make sure server is stopped
        rt.httpsvr.ul_svr_wait(s)
        return rc

    def cp_from(self, rpath, lpath):
        print("COPY FROM: %s -> %s" % (rpath, lpath))
        if not rpath.endswith('/'):
            self._cp_from(rpath, lpath)

        # local temporary file (only use it for the filename)
        fd, lfp = tempfile.mkstemp(suffix=".tar.gz")
        os.close(fd)
        _, fn = os.path.split(lfp)

        # we have a directory to copy.
        # 1. Tar up directory to a temporary file
        rfn = os.path.join("/c/Users/%s/AppData/Local/" % self.user, fn)
        print(rfn)
        # get parent directory and directory to tar
        rc = self.run_sh("(cd %s && tar czf %s .)" %
                         (rpath, rfn))
        if not rc == 0:
            return rc

        # 2. Copy file to the host and delete the remote file
        rc = self._cp_from(rfn, lpath)
        self.run_sh("rm -f %s" % rfn)
        if not rc == 0:
            return rc

        # 3. Untar and remove local tar file
        rc = subprocess.call("(cd %s; tar xzf %s)" % (lpath, fn), shell=True)
        os.remove(os.path.join(lpath, fn))
        return rc

    def run_rt_local(self, rpath, rbin, args):
        # When running tests pipe output to /dev/null. We can only get
        # it after the command was run anyway so it's a bit
        # useless. Instead we run a Logging server in a thread and
        # configure rt-local to log to it.
        bg = False
        if " run" in args:
            thd = _LogThread(_PORT_AUX_LOG)
            thd.start()
            args = "--logger %s:%d " % (self.our_ip, _PORT_AUX_LOG) + args
            bg = True

        # base64 encode the rt-local arguments in case of special chars
        rc = _Host.run_rt_local(self, rpath, rbin, "%s %s" %
                                (base64.b64encode(args),
                                 " > /dev/null 2> /dev/null &" if bg else ""))

        if " run" in args:
            print("'rt-local' started in background on remote host")
            thd.join()
        # XXX ToDo: reflect the rt-local return value here
        return rc


class SSHHost(_Host):
    """A class to interact with remote host via SSH"""

    def __init__(self, remote, user, password):
        super(SSHHost, self).__init__(remote, user, password)
        self.connection_string = "%s@%s" % (self.user, self.remote)
        self.key_file = _KEYFILE
        self.ignore_host_key = '-oStrictHostKeyChecking=no'
        self.known_hosts_file = '-oUserKnownHostsFile=/dev/null'

        # Check SSH key has correcy permissions
        perms = stat.S_IRUSR & stat.S_IWUSR
        st = os.stat(self.key_file)
        current_perms = stat.S_IMODE(st.st_mode)
        if current_perms != perms:
            os.chmod(self.key_file, 0o600)

    def run_sh(self, cmd):
        """Run a shell command on the host"""
        print("EXEC: %s" % cmd)
        ssh_cmd = [
            "ssh",
            "-q",
            self.ignore_host_key,
            self.known_hosts_file,
            "-i",
            self.key_file,
            self.connection_string,
            '%s' % (cmd)
        ]
        # Use Popen to avoid jamming the stdout buffer
        process = subprocess.Popen(ssh_cmd,
                                   stdin=None,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.STDOUT)
        for line in iter(process.stdout.readline, ''):
            sys.stdout.write(line)

        process.communicate()
        return process.returncode

    def cp_from(self, rpath, lpath):
        if not rpath.endswith('/'):
            rc = self.scp_from(False, rpath, lpath)
        else:
            rc = self.scp_from(True, rpath, lpath)
        return rc

    def _cp_to(self, lpath, rpath):
        scp_cmd = [
            "scp",
            "-q",
            self.ignore_host_key,
            "-i",
            self.key_file,
            lpath,
            "%s:%s" % (self.connection_string, rpath)
        ]
        code = subprocess.call(scp_cmd)
        return code

    def scp_from(self, is_dir, rpath, lpath):
        print("COPY FROM: %s -> %s" % (rpath, lpath))
        if is_dir:
            scp = ["scp", "-q", "-r"]
            if not rpath.endswith("*"):
                rpath = rpath + "*"
        else:
            scp = ["scp", "-q"]

        scp_cmd = scp + [
            self.ignore_host_key,
            "-i",
            self.key_file,
            "%s:%s" % (self.connection_string, rpath),
            lpath
        ]
        code = subprocess.call(scp_cmd)
        return code


def main():
    """Main entry point to run regressions tests on a (remote) host"""

    # Make sure stdout is not buffered
    sys.stdout = Unbuffered(sys.stdout)

    parser = argparse.ArgumentParser()
    parser.description = \
        "Execute rt-local commands on a (remote) host"
    parser.epilog = \
        "Arguments after '--' are passed to rt-local. " + \
        "If no '--' is provided only the tests are copied " + \
        "over to the remote system."

    parser.add_argument('-t', '--type', nargs=1, required=True, dest='rtype',
                        choices=['osx', 'win'],
                        help="Type of host")

    parser.add_argument("-r", "--remote", required=True,
                        help="Remote host IP or name")

    parser.add_argument("--portbase", type=int, default=12000,
                        help="Port(s) to use for auxiliary services")

    parser.add_argument("--local-ip", type=str, dest='localip',
                        help="Override local IP clients connect to.")

    parser.add_argument("-u", "--user", required=True,
                        help="Username to use")

    parser.add_argument("-p", "--password", required=True,
                        help="Password for user")

    parser.add_argument("-f", "--file", type=str, dest='file',
                        help="File to upload to RT_ROOT")

    parser.add_argument("--results", action='store_true',
                        help="Copy results from the remote host")

    parser.add_argument("-c", "--cases", type=str, dest='cases',
                        default="./",
                        help="Folder containing test cases")
    parser.add_argument("-k", "--keyfile", type=str, dest='keyfile',
                        help="SSH key file")

    # Split argv at -- and parse
    host_args, _, local_args = " ".join(sys.argv).partition(" -- ")
    args = parser.parse_args(host_args.split()[1:])

    global _PORT_AUX_HTTP
    global _PORT_AUX_LOG
    global _LOCAL_IP
    _PORT_AUX_HTTP = args.portbase
    _PORT_AUX_LOG = args.portbase + 1
    _LOCAL_IP = args.localip
    _KEYFILE = args.keyfile

    if args.rtype[0] == "osx":
        if _KEYFILE is None:
            print("ERROR: Required argument --keyfile not supplied")
            return 1
        h = SSHHost(args.remote, args.user, args.password)
        remote_tmp_dir = _OSX_TMP_DIR
        rt_install_cmd = "sudo "
        rt_path = "/usr/local/bin/rt-local"
    elif args.rtype[0] == "win":
        h = WinRMHost(args.remote, args.user, args.password)
        remote_tmp_dir = _WIN_TMP_DIR % args.user
        rt_install_cmd = ""
        rt_path = "rt-local"

    if args.results:
        print("\n=== Copy Results")
        return h.cp_from(remote_tmp_dir + "/_results/", "./_results")

    # prepare system
    print("\n=== Clean Remote system")
    h.run_sh("rm -rf %s" % remote_tmp_dir)
    h.run_sh("mkdir %s" % remote_tmp_dir)

    # prepare rt-tools
    print("\n=== Install rt-framework on remote host")
    base_tools_dir = tempfile.mkdtemp()
    tools_dir = os.path.join(base_tools_dir, "rt-tools")
    os.mkdir(tools_dir)

    rc = subprocess.call(
        "easy_install -qzmaxld %s rt-framework" % tools_dir,
        stdout=DEVNULL,
        stderr=DEVNULL,
        shell=True
    )
    if rc > 0:
        print("Error installing framework: %d" % rc)
        return rc

    rc = h.cp_to(tools_dir, remote_tmp_dir)
    shutil.rmtree(base_tools_dir)
    if rc > 0:
        print("Error copying framework: %d" % rc)
        return rc

    # Install the rt tools on the remote system
    rt_install_cmd += (
        'find %s -iname "*.egg" -exec easy_install -qN -H None '
        '-f $(dirname {}) $(basename {}) \;'
        % (remote_tmp_dir + "/rt-tools")
    )
    rc = h.run_sh(rt_install_cmd)
    if rc > 0:
        print("Error installing framework remotely: %d" % rc)
        return rc

    print("\n=== Copy infrastructure and test cases")
    rc = h.cp_to(args.cases, remote_tmp_dir, ["./_results/*",
                                              "*.pyc",
                                              "./cases/_tmp/*",
                                              "./etc/*"])
    if rc > 0:
        print("Error copying tests: %d" % rc)
        return rc

    if args.file is not None:
        h.cp_to(args.file, remote_tmp_dir)

    if len(local_args) == 0:
        return

    print("\n=== Run 'rt-local' on remote host")
    rc = h.run_rt_local(remote_tmp_dir, rt_path, local_args)

    if " run" in local_args:
        print("\n=== Copy Results")
        t = h.cp_from(remote_tmp_dir + "/_results/", "./_results")
        rc = t if rc == 0 else rc

    return rc