#!/usr/bin/env python3
# -*- 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.

"""


import contextlib
import logging
import time
from optparse import OptionParser
import os
import shutil
import subprocess
from subprocess import CalledProcessError, Popen, PIPE, STDOUT
import sys
import tempfile
import urllib.request, urllib.error
from shlex import quote


OSG_CA_SCRIPTS_REPO = "https://github.com/opensciencegrid/osg-ca-scripts"
OSG_CA_SCRIPTS_BRANCH = "latest"


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(osg_ca_scripts_dir, cert_dir):
    """Run osg-ca-manage setupCA from a git clone of osg-ca-scripts.

    This will write CA certificates into cert_dir.  As a side effect, it will
    also create a symlink in osg_ca_scripts_dir/etc/grid-security/certificates
    but we can ignore that.

    """
    if os.path.basename(cert_dir) != "certificates":  # required by osg-ca-manage
        raise ValueError("cert_dir %r does not end in certificates" % cert_dir)
    environ = os.environ.copy()
    environ["OSG_LOCATION"] = osg_ca_scripts_dir
    environ["PERL5LIB"] = "%s:%s/lib" % (environ.get("PERL5LIB", ""), osg_ca_scripts_dir)

    # 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 = ["./sbin/osg-ca-manage"]
    command += ["setupCA"]
    command += ["--location", os.path.dirname(cert_dir)]
    command += ["--url", "osg"]
    with chdir(osg_ca_scripts_dir):
        subprocess.check_call(command, env=environ)


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 (b"CRL verification failed" in output or b"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:
        # Suppress any SSH pre-login banners 
        ssh.extend(["-q"])
        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
    
    # Suppress any SSH pre-login banners 
    ssh.extend(["-q"])

    errstr = "Error rsyncing to remote host %s:%s: " % (remote_host, remote_dir)
    try:
        proc = Popen(
            ssh + [remote_host, "[[ -e %s ]] || echo missing" % quote(remote_dir)],
            stdout=PIPE,
        )
    except OSError as e:
        raise Error(errstr + str(e))
    output, _ = proc.communicate()
    output = output.decode("latin-1")
    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(
                quote(olddir), quote(remote_dir), quote(newdir))])
    except (OSError, CalledProcessError) as e:
        raise Error("Error renaming remote directories: %s" % e)


def setup_osg_ca_scripts_dir(osg_ca_scripts_dir):
    """Clones the osg-ca-scripts repo and sets up dirs and symlinks to allow
    running the script right from the git checkout.

    """
    retries = 3
    while retries > 0:
        ret = subprocess.call([
            "git", "clone", "-b", OSG_CA_SCRIPTS_BRANCH, OSG_CA_SCRIPTS_REPO, osg_ca_scripts_dir
        ])
        if ret == 0:
            break
        retries -= 1
        if retries > 0:
            log.warning("git clone failed; retrying in 30 seconds")
            time.sleep(30)
        else:
            raise Error("git clone failed")
    # osg-ca-manage can be run right from the git checkout with a few tweaks:
    # - Need to set $OSG_LOCATION to the root of the checkout (handled in setup_cas())
    # - Config file must be in $OSG_LOCATION/etc/osg/osg-update-certs.conf
    # - osg-setup-ca-certificates must be in $OSG_LOCATION/usr/libexec
    # - OSGCerts.pem must be in Perl path (handled in setup_cas())
    with chdir(osg_ca_scripts_dir):
        os.symlink(".", "usr")
        os.mkdir("etc/osg")
        os.symlink("../osg-update-certs.conf", "etc/osg/osg-update-certs.conf")


@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)


@contextlib.contextmanager
def chdir(directory):
    """Resource manager for cd'ing into an existing directory
    and going back to the old directory afterward.

    """
    olddir = os.getcwd()
    os.chdir(directory)
    yield
    os.chdir(olddir)


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")
    parser.add_option("--dry-run", action="store_true", help="Do not deploy to remote host")
    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 opts.dry_run:
        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")

            osg_ca_scripts_dir = os.path.join(wd, "osg-ca-scripts")
            setup_osg_ca_scripts_dir(osg_ca_scripts_dir)

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

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

            if opts.dry_run:
                log.info("Not uploading; would run:")
                log.info("rsync_upload(%r, %r, %r, %r, %r)" % (deploy_client_dir, opts.remote_user, remote_host, opts.remote_dir, opts.ssh_key))
            else:
                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())
