#!/usr/bin/env python2
# -*- coding: utf-8 -*-
# Requirements: Python 2.6+
"""Create or update a tarball-based worker node client installation on a
remote host from a hosted CE.

This downloads the worker node tarball and creates an installation in a
temporary directory into which CAs (from the OSG CA distribution) and CRLs
are downloaded.  Then the installation is uploaded using rsync to the remote
host.

SSH access to the remote host and write access to the destination directory
is required.  The parent of the destination directory must already exist on
the remote host.

"""

from __future__ import print_function
import contextlib
import logging
from optparse import OptionParser
import os
import shutil
import subprocess
from subprocess import CalledProcessError, Popen, PIPE, STDOUT
import sys
import tempfile

from six.moves import shlex_quote
from six.moves import urllib


devnull = open(os.devnull, "w+")
log = logging.getLogger(__name__)


class Error(Exception):
    pass


# adapted from osgbuild/fetch_sources; thanks Carl
def download_to_file(uri, outfile):
    try:
        handle = urllib.request.urlopen(uri)
    except urllib.error.URLError as err:
        raise Error("Unable to download %s: %s" % (uri, err))

    try:
        with open(outfile, "wb") as desthandle:
            chunksize = 64 * 1024
            chunk = handle.read(chunksize)
            while chunk:
                desthandle.write(chunk)
                chunk = handle.read(chunksize)
    except EnvironmentError as err:
        raise Error("Unable to save downloaded file to %s: %s" % (outfile, err))


def setup_cas(wn_client, osgrun, cert_dir):
    "Run osg-ca-manage setupCA"

    # Unfortunately, if we specify a location to osg-ca-manage setupCA, it
    # always wants to create a symlink from
    # $OSG_LOCATION/etc/grid-security/certificates to that location.  Since we
    # do not want to mess up the tarball install we're using, we must first
    # save the symlink that's already there, then run osg-ca-manage, then
    # restore it.
    certs_link = os.path.join(wn_client, "etc/grid-security/certificates")
    certs_link_save = certs_link + ".save"

    # Note that in a proper tarball install, certs_link should already exist
    # but handle its nonexistence gracefully too.
    # Note the need to use 'lexists' since 'exists' returns False if the path
    # is a broken symlink.
    if os.path.lexists(certs_link):
        if os.path.lexists(certs_link_save):
            os.unlink(certs_link_save)
        os.rename(certs_link, certs_link_save)

    # osg-ca-manage always puts the certs into a subdirectory called 'certificates'
    # under the location specified here. So specify the parent of cert_dir as --location.
    command = [osgrun, "osg-ca-manage"]
    command += ["setupCA"]
    command += ["--location", os.path.dirname(cert_dir)]
    command += ["--url", "osg"]
    try:
        subprocess.check_call(command)
    finally:
        if os.path.lexists(certs_link_save):
            if os.path.lexists(certs_link):
                os.unlink(certs_link)
            os.rename(certs_link_save, certs_link)


def update_cas(osgrun, cert_dir):
    subprocess.check_call([osgrun, "osg-ca-manage", "--cert-dir", cert_dir, "refreshCA"])


def update_crls(cert_dir):
    """Run system fetch-crl; ignore non-fatal errors, raise on others.

    Run the system fetch-crl instead of fetch-crl in the tarball install
    because fetch-crl must be compiled for the same OS as the script runs on.
    """
    command = ["fetch-crl"]
    command += ["--infodir", cert_dir]
    command += ["--out", cert_dir]
    command += ["--quiet"]
    command += ["--agingtolerance", "24"]  # 24 hours
    command += ["--parallelism", "5"]

    output = None
    proc = Popen(command, stdout=PIPE, stderr=STDOUT)
    output, _ = proc.communicate()
    if proc.returncode != 0:
        if output and ("CRL verification failed" in output or "Download error" in output):
            # These errors aren't actually fatal; we'll send a less alarming
            # notification about them.
            log.info(output)
        else:
            log.error(output)
            raise Error("fetch-crl failed with error code %d" % proc.returncode)


def check_connectivity(remote_user, remote_host, ssh_key):
    ssh = ["ssh"]
    if remote_user:
        ssh.extend(["-l", remote_user])
    if ssh_key:
        ssh.extend(["-i", ssh_key])
    try:
        subprocess.check_call(ssh + [remote_host, "true"])
    except CalledProcessError:
        return False
    return True


def rsync_upload(local_dir, remote_user, remote_host, remote_dir, ssh_key=None):
    # type: (str, str, str, str, str) -> None
    """Use rsync to upload the contents of a directory to a remote host,
    minimizing the time the remote dir spends in an inconsistent state.
    Requires rsync and ssh shell access on the remote host to do the swapping.

    The parent directories must already exist.
    """
    ssh = ["ssh"]
    if remote_user:
        ssh.extend(["-l", remote_user])
    if ssh_key:
        ssh.extend(["-i", ssh_key])
    olddir = "%s~old~" % remote_dir
    newdir = "%s~new~" % remote_dir
    local_dir = local_dir.rstrip("/") + "/"  # exactly 1 trailing slash

    errstr = "Error rsyncing to remote host %s:%s: " % (remote_host, remote_dir)
    try:
        proc = Popen(
            ssh + [remote_host, "[[ -e %s ]] || echo missing" % shlex_quote(remote_dir)],
            stdout=PIPE,
        )
    except OSError as e:
        raise Error(errstr + str(e))
    output, _ = proc.communicate()
    if proc.returncode != 0:
        log.error(output)
        raise Error(errstr + "rsync exited with %d" % proc.returncode)

    try:
        if output.rstrip() == "missing":
            log.info("rsyncing entire WN client to %s:%s", remote_host, remote_dir)
            # If remote dir is missing then just upload and return
            subprocess.check_call(["rsync", "-e", " ".join(ssh),
                                   "-qaz",
                                   local_dir,
                                   "%s:%s" % (remote_host, remote_dir)])
            return

        # Otherwise, upload to newdir
        log.info("rsyncing WN client changes to %s:%s", remote_host, newdir)
        subprocess.check_call(["rsync", "-e", " ".join(ssh),
                               "-qaz",
                               "--link-dest", remote_dir,
                               "--delete-before",
                               local_dir,
                               "%s:%s" % (remote_host, newdir)])
    except (OSError, CalledProcessError) as e:
        raise Error(errstr + str(e))

    # then rename destdir to olddir and newdir to destdir
    try:
        log.info("Moving %s to %s", newdir, remote_dir)
        subprocess.check_call(ssh +
                              [remote_host,
             "rm -rf {0} && "
             "mv {1} {0} && "
             "mv {2} {1}".format(
                shlex_quote(olddir), shlex_quote(remote_dir), shlex_quote(newdir))])
    except (OSError, CalledProcessError) as e:
        raise Error("Error renaming remote directories: %s" % e)


@contextlib.contextmanager
def working_dir(*args, **kwargs):
    """Resource manager for creating a temporary directory, cd'ing into it,
    and deleting it after completion.

    """
    wd = tempfile.mkdtemp(*args, **kwargs)
    olddir = os.getcwd()
    os.chdir(wd)
    yield wd
    os.chdir(olddir)
    shutil.rmtree(wd)


# Because paths are embedded into the WN-client installation, two copies
# of the client are used: the "deploy" client, which will be rsynced to the
# worker node, and the "fetch" client, which will be used to download CAs
# and CRLs.  The "fetch" client will put the CAs and CRLs into the certificate
# dir of the "deploy" client.


def main():
    parser = OptionParser(usage="usage: %prog [options] remote_host", description=__doc__)
    parser.add_option(
        "--upstream-url",
        default="https://repo.opensciencegrid.org/tarball-install/3.4/osg-wn-client-latest.el7.x86_64.tar.gz",
        help="URL for the WN tarball file. [default: %default]",
    )
    parser.add_option("--remote-user", help="remote user to use for rsync and ssh")
    parser.add_option(
        "--remote-dir",
        default="/home/bosco/osg-wn-client",
        help="remote directory the WN client will be placed in. [default: %default]",
    )
    parser.add_option("--ssh-key", help="SSH key to use to log in with")
    opts, args = parser.parse_args()
    if len(args) != 1:
        parser.error("incorrect number of arguments")
    remote_host = args[0]

    # check if rsync is installed and working
    try:
        subprocess.check_call(["rsync", "--version"], stdout=devnull)
    except (CalledProcessError, EnvironmentError) as e:
        log.error("Error invoking rsync: %s", e)
        return 1

    if not check_connectivity(opts.remote_user, remote_host, opts.ssh_key):
        log.error("Could not connect to remote host")
        return 1

    with working_dir() as wd:
        try:
            log.info("Downloading WN tarball")
            download_to_file(opts.upstream_url, "osg-wn-client.tar.gz")

            os.mkdir("deploy")
            subprocess.check_call(["tar", "-C", "deploy", "-xzf", "osg-wn-client.tar.gz"])
            deploy_client_dir = os.path.join(wd, "deploy/osg-wn-client")
            cert_dir = os.path.join(deploy_client_dir, "etc/grid-security/certificates")

            os.mkdir("fetch")
            subprocess.check_call(["tar", "-C", "fetch", "-xzf", "osg-wn-client.tar.gz"])
            fetch_client_dir = os.path.join(wd, "fetch/osg-wn-client")

            log.info("Setting up tarball dirs")
            subprocess.check_call([os.path.join(deploy_client_dir, "osg/osg-post-install"),
                                   "-f", opts.remote_dir])
            subprocess.check_call([os.path.join(fetch_client_dir, "osg/osg-post-install")])
            osgrun = os.path.join(fetch_client_dir, "osgrun")

            log.info("Fetching CAs")
            setup_cas(fetch_client_dir, osgrun, cert_dir)
            log.info("Fetching CRLs")
            update_crls(cert_dir)

            log.info("Uploading")
            rsync_upload(deploy_client_dir, opts.remote_user, remote_host, opts.remote_dir, opts.ssh_key)
        except (EnvironmentError, CalledProcessError, Error) as e:
            log.error(e)
            return 1

    return 0


if __name__ == "__main__":
    logging.basicConfig(format="*** %(message)s", level=logging.INFO)
    sys.exit(main())
