image-download

#!/usr/bin/env python3

# This file is part of Cockpit.
#
# Copyright (C) 2013 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

#
# Download images or other state
#
# Images usually have a name specific link committed to git. These
# are referred to as 'committed'
#
# Other state is simply referenced by name without a link in git
# This is referred to as 'state'
#
# The stores are places to look for images or other state
#

import argparse
import email
import io
import os
import shutil
import socket
import stat
import subprocess
import sys
import tempfile
import time
import fcntl
import urllib.parse

from machine.machine_core.constants import IMAGES_DIR
from machine.machine_core.directories import get_images_data_dir

from task import PUBLIC_STORES, REDHAT_STORES
from task.testmap import get_test_image

CONFIG = "~/.config/image-stores"

DEVNULL = open("/dev/null", "r+")
EPOCH = "Thu, 1 Jan 1970 00:00:00 GMT"


def find(name, stores, latest, quiet):
    found = []
    ca = os.path.join(IMAGES_DIR, "files", "ca.pem")

    for store in stores:
        url = urllib.parse.urlparse(store)

        defport = url.scheme == 'http' and 80 or 443

        try:
            ai = socket.getaddrinfo(url.hostname, url.port or defport, socket.AF_INET, 0, socket.IPPROTO_TCP)
        except socket.gaierror:
            ai = []
            message = store

        for (family, socktype, proto, canonname, sockaddr) in ai:
            message = "{scheme}://{0}:{1}{path}".format(*sockaddr, scheme=url.scheme, path=url.path)

            def curl(*args):
                try:
                    cmd = ["curl"] + list(args) + ["--head", "--connect-timeout", "10", "--silent",
                                                   "--fail", "--cacert", ca, source]
                    start = time.time()
                    output = subprocess.check_output(cmd, universal_newlines=True)
                    found.append((cmd, output, message, time.time() - start))
                    if not quiet:
                        sys.stderr.write(" > {0}\n".format(message))
                    return True
                except subprocess.CalledProcessError:
                    return False

            # first try with stores that accept the "cockpit-tests" host name
            resolve = "cockpit-tests:{1}:{0}".format(*sockaddr)
            source = urllib.parse.urljoin("{0}://cockpit-tests:{1}{2}".format(url.scheme, sockaddr[1], url.path), name)
            if curl("--resolve", resolve):
                continue

            # fall back for OpenShift proxied stores which send their own SSL cert initially
            # host name has to match that
            source = urllib.parse.urljoin(store, name)
            if curl():
                continue

            if not quiet:
                sys.stderr.write(" x {0}\n".format(message))

    # If we couldn't find the file, but it exists, we're good
    if not found:
        return None, None

    # Find the most recent version of this file
    def header_date(args):
        cmd, output, message, latency = args
        try:
            reply_line, headers_alone = output.split('\n', 1)
            last_modified = email.message_from_file(io.StringIO(headers_alone)).get("Last-Modified", "")
            return time.mktime(time.strptime(last_modified, '%a, %d %b %Y %H:%M:%S %Z'))
        except ValueError:
            return ""

    if latest:
        found.sort(reverse=True, key=header_date)
    else:
        found.sort(reverse=False, key=lambda x: x[3])

    # Return the command and message
    return found[0][0], found[0][2]


def download(dest, force, state, quiet, stores):
    if not stores:
        config = os.path.expanduser(CONFIG)
        if os.path.exists(config):
            with open(config, 'r') as fp:
                stores = fp.read().strip().split("\n")
        else:
            stores = []
        stores += PUBLIC_STORES
        stores += REDHAT_STORES

    # The time condition for If-Modified-Since
    exists = not force and os.path.exists(dest)
    if exists:
        since = dest
    else:
        since = EPOCH

    name = os.path.basename(dest)
    cmd, message = find(name, stores, latest=state, quiet=quiet)

    # If we couldn't find the file, but it exists, we're good
    if not cmd:
        if exists:
            return
        raise RuntimeError("image-download: couldn't find file anywhere: {0}".format(name))

    # Choose the first found item after sorting by date
    if not quiet:
        sys.stderr.write(" > {0}\n".format(urllib.parse.urljoin(message, name)))

    temp = dest + ".partial"

    # Adjust the command above that worked to make it visible and download real stuff
    cmd.remove("--head")
    cmd.append("--show-error")
    if not quiet and os.isatty(sys.stdout.fileno()):
        cmd.remove("--silent")
        cmd.insert(1, "--progress-bar")
    cmd.append("--remote-time")
    cmd.append("--time-cond")
    cmd.append(since)
    cmd.append("--output")
    cmd.append(temp)
    if os.path.exists(temp):
        if force:
            os.remove(temp)
        else:
            cmd.append("-C")
            cmd.append("-")

    # Always create the destination file (because --state)
    else:
        open(temp, 'a').close()

    curl = subprocess.Popen(cmd)
    ret = curl.wait()
    if ret != 0:
        raise RuntimeError("curl: unable to download %s (returned: %s)" % (message, ret))

    os.chmod(temp, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)

    # Due to time-cond the file size may be zero
    # A new file downloaded, put it in place
    if not exists or os.path.getsize(temp) > 0:
        shutil.move(temp, dest)


# Calculate a place to put images where links are not committed in git
def state_target(path):
    data_dir = get_images_data_dir()
    os.makedirs(data_dir, mode=0o775, exist_ok=True)
    return os.path.join(data_dir, path)


# Calculate a place to put images where links are committed in git
def committed_target(image):
    link = os.path.join(IMAGES_DIR, image)
    if not os.path.islink(link):
        raise RuntimeError("image link does not exist: " + image)

    dest = os.readlink(link)
    relative_dir = os.path.dirname(os.path.abspath(link))
    full_dest = os.path.join(relative_dir, dest)
    while os.path.islink(full_dest):
        link = full_dest
        dest = os.readlink(link)
        relative_dir = os.path.dirname(os.path.abspath(link))
        full_dest = os.path.join(relative_dir, dest)

    dest = os.path.join(get_images_data_dir(), dest)

    # We have the file but there is not valid link
    if os.path.exists(dest):
        try:
            os.symlink(dest, os.path.join(IMAGES_DIR, os.readlink(link)))
        except FileExistsError:
            pass

    # The image file in the images directory, may be same as dest
    image_file = os.path.join(IMAGES_DIR, os.readlink(link))

    # Double check that symlink in place but never make a cycle.
    if os.path.abspath(dest) != os.path.abspath(image_file):
        try:
            os.symlink(os.path.abspath(dest), image_file)
        except FileExistsError:
            pass

    return dest


def wait_lock(target):
    lockfile = os.path.join(tempfile.gettempdir(), ".cockpit-test-resources", os.path.basename(target) + ".lock")
    os.makedirs(os.path.dirname(lockfile), exist_ok=True)

    # we need to keep the lock fd open throughout the entire runtime, so remember it in a global-scoped variable
    wait_lock.f = open(lockfile, "w")
    for retry in range(360):
        try:
            fcntl.flock(wait_lock.f, fcntl.LOCK_NB | fcntl.LOCK_EX)
            return
        except BlockingIOError:
            if retry == 0:
                print("Waiting for concurrent image-download of %s..." % os.path.basename(target))
            time.sleep(10)
    else:
        raise TimeoutError("timed out waiting for concurrent downloads of %s\n" % target)


def download_images(image_list, force, quiet, state, store):
    data_dir = get_images_data_dir()
    os.makedirs(data_dir, exist_ok=True)

    # A default set of images are all links in git.  These links have
    # no directory part.  Other links might exist, such as the
    # auxiliary links created by committed_target above, and we ignore
    # them.
    if not image_list:
        image_list = []
        if not state:
            for filename in os.listdir(IMAGES_DIR):
                link = os.path.join(IMAGES_DIR, filename)
                if os.path.islink(link) and os.path.dirname(os.readlink(link)) == "":
                    image_list.append(filename)

    success = True

    for image in image_list:
        image = get_test_image(image)
        try:
            if state:
                target = state_target(image)
            else:
                target = committed_target(image)

            # don't download the same thing multiple times in parallel
            wait_lock(target)

            if force or state or not os.path.exists(target):
                download(target, force, state, quiet, store)
        except Exception as ex:
            success = False
            sys.stderr.write("image-download: {0}\n".format(str(ex)))

    return success


def main():
    parser = argparse.ArgumentParser(description='Download a bot state or images')
    parser.add_argument("--force", action="store_true", help="Force unnecessary downloads")
    parser.add_argument("--store", action="append", help="Where to find state or images")
    parser.add_argument("--quiet", action="store_true", help="Make downloading quieter")
    parser.add_argument("--state", action="store_true", help="Images or state not recorded in git")
    parser.add_argument('image', nargs='*')
    args = parser.parse_args()

    if not download_images(args.image, args.force, args.quiet, args.state, args.store):
        return 1

    return 0


if __name__ == '__main__':
    sys.exit(main())