cvmfs-singularity-sync

#!/usr/bin/env python3

# NOTE:
# This file is forked from singularity 2.2's "cli.py".  I have left the
# copyright notice below untouched; this file remains under the same license.
# - Brian Bockelman

'''

bootstrap.py: python helper for Singularity command line tool

Copyright (c) 2016, Vanessa Sochat. All rights reserved.

"Singularity" Copyright (c) 2016, The Regents of the University of California,
through Lawrence Berkeley National Laboratory (subject to receipt of any
required approvals from the U.S. Dept. of Energy).  All rights reserved.

This software is licensed under a customized 3-clause BSD license.  Please
consult LICENSE file distributed with the sources of this project regarding
your rights to use or distribute this software.

NOTICE.  This Software was developed under funding from the U.S. Department of
Energy and the U.S. Government consequently retains certain rights. As such,
the U.S. Government has been granted for itself and others acting on its
behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software
to reproduce, distribute copies to the public, prepare derivative works, and
perform publicly and display publicly, and to permit other to do so.


'''

import sys
import argparse
import os
import errno
import fnmatch
import json
import urllib.request, urllib.error, urllib.parse
import hashlib
import traceback
import subprocess
import dockerhub
import cleanup
import sqlitedict
import glob

DOCKER_CREDS = dict()

def main():
    parser = argparse.ArgumentParser(description="Bootstrap Docker images for Singularity containers deployed to CVMFS")

    # Name of the docker image, required
    parser.add_argument("--docker",
                        dest='docker',
                        help="name of Docker image to bootstrap, in format library/ubuntu:latest",
                        type=str,
                        default=None)

    # Link to the file of the docker images, required
    parser.add_argument("--images-url",
                        dest='filelist',
                        help="URL to download a list of Docker images from",
                        type=str,
                        default=None)

    # Name of the docker image file, required
    parser.add_argument("--images-file",
                        dest='filelist_path',
                        help="Local file path to a list of Docker images",
                        type=str,
                        default=None)

    # root file system of singularity image
    parser.add_argument("--rootfs",
                        dest='rootfs',
                        help="the path for the root filesystem to extract to",
                        type=str,
                        default=None)

    # Manifest cache file
    parser.add_argument("--sqlite-cache",
                        help="SQLite database to cache Docker manifests",
                        type=str,
                        default=None)

    # Docker registry (default is registry-1.docker.io)
    parser.add_argument("--registry",
                        dest='registry',
                        help="the registry path to use, to replace registry-1.docker.io",
                        type=str,
                        default=None)

    # Flag to indicate a token is not required
    parser.add_argument("--no-token",
                        dest='notoken',
                        action="store_true",
                        help="Indicate that auth is not required",
                        default=False)

    # File with access token
    parser.add_argument("--docker-creds",
                        help="JSON file with Docker credentials",
                        type=argparse.FileType('r'))

    # Flag to indicate dry-run
    parser.add_argument("--dry-run",
                        dest='dryrun',
                        action="store_true",
                        help="Indicate that this is a dry-run",
                        default=False)

    try:
        args = parser.parse_args()
    except:
        parser.print_help()
        sys.exit(0)

    if args.docker_creds:
        global DOCKER_CREDS
        DOCKER_CREDS = json.load(args.docker_creds)
        args.docker_creds.close()

    manifest_cache = sqlitedict.SqliteDict(args.sqlite_cache, tablename='manifests', autocommit=True)

    # Find root filesystem location
    if args.rootfs:
        singularity_rootfs = args.rootfs
    else:
        singularity_rootfs = '/cvmfs/singularity.opensciencegrid.org'
    singularity_rootfs = os.path.abspath(singularity_rootfs)

    # Does the registry require a token?
    doauth = not args.notoken

    # Do we have a docker image specified?
    if not args.docker and not (args.filelist or args.filelist_path):
        print("No docker image or file list specified..", file=sys.stderr)
        return 1

    currentImages = []
    if args.docker:
        image = args.docker
        if not args.dryrun:
            return publish_image(image, singularity_rootfs, args.registry, doauth, manifest_cache)
        else:
            return verify_image(image, args.registry, doauth, manifest_cache)
    else:
        final_retval = 0
        failed_images = []
        #fp = urllib.request.urlopen(args.filelist).read().decode().split('\n')
        if args.filelist:
            fp = urllib.request.urlopen(args.filelist).read().decode().split('\n')
        else:
            if args.filelist_path:
                fp = open(args.filelist_path, "r").read().split('\n')
        for image in fp:
            image = image.strip()
            if image.startswith("#") or not image:
                continue
            registry, namespace, repo_name, repo_tag = parse_image(image)
            print("Working on image: {}".format(image))

            if '*' in repo_tag: # Treat wildcards as a glob
                try:
                    tag_names = get_tags(namespace, repo_name, registry=registry, auth=doauth)
                except Exception as ex:
                    image = '%s/%s/%s' % (registry, namespace, repo_name)
                    print("Failed to get tags for image: {}".format(image))
                    traceback.print_exc()
                    print(ex)
                    glob_path = os.path.join(singularity_rootfs, namespace, repo_name + ":" + repo_tag)
                    existing_images = glob.glob(glob_path)

                    # Remove the prepended singularity_rootfs and slash
                    existing_images = [x[len(singularity_rootfs)+1:] for x in existing_images]
                    for existing_image in existing_images:
                        registry, namespace, repo_name, repo_tag = parse_image(existing_image)
                        currentImages.append('%s/%s/%s:%s' % (registry, namespace, repo_name, repo_tag))

                    continue
                repo_tag = fnmatch.filter(tag_names, repo_tag)
            else:
                repo_tag = [repo_tag]

            for tag in repo_tag:
                image = '%s/%s/%s:%s' % (registry, namespace, repo_name, tag)
                currentImages.append(image)
                retval = 1
                tries = 3
                for i in range(tries):
                    if not args.dryrun:
                        try:
                            retval = publish_image(image, singularity_rootfs, registry, doauth, manifest_cache)
                        except Exception as ex:
               	            if i < tries -1:
                                print("Failed to publish image: {}".format(image))
                                traceback.print_exc()
                                print("Retry image {}".format(image))
                                continue
                            else:
                                print("Tried {} times ".format(tries) + "for image {}".format(image) + ", giving up")
                    else:
                        try:
                            retval = verify_image(image, registry, doauth, manifest_cache)
                        except Exception as ex:
                            if i < tries -1:
                                print("Failed to verify image: {}".format(image))
                                traceback.print_exc()
                                print("Retry image {}".format(image))
                                continue
                            else:
                                print("Tried {} times ".format(tries) + "for image {}".format(image))
                    break
                if retval:
                    final_retval = retval
                    failed_images.append("{}".format(image))
        print("All requested images have been attempted; final return code: %d" % final_retval)
        if final_retval:
            print("Failed images:")
            print(*failed_images, sep = "\n")
        if not args.dryrun:
            print("Cleaning up unlinked images")
            start_txn(singularity_rootfs)
            cleanup.remove_unlisted_images(currentImages, singularity_rootfs)
            cleanup.cleanup(singularity_base=singularity_rootfs)
            publish_txn()
        return final_retval


_in_txn = False
def start_txn(singularity_rootfs):
    global _in_txn
    if _in_txn:
        return 0
    if not singularity_rootfs.startswith("/cvmfs/singularity.opensciencegrid.org"):
        return 0
    if os.path.exists("/var/spool/cvmfs/singularity.opensciencegrid.org/in_transaction.lock"):
        result = os.system("cvmfs_server abort -f singularity.opensciencegrid.org")
        if result:
            print("Failed to abort lingering transaction (exit status %d)." % result, file=sys.stderr)
            return 1
    result = os.system("cvmfs_server transaction singularity.opensciencegrid.org")
    if result:
        print("Transaction start failed (exit status %d); will not attempt update." % result, file=sys.stderr)
        return 1
    _in_txn = True

    # Test CVMFS mount if applicable.
    test_dir = os.path.join(singularity_rootfs, "library")
    if not os.path.exists(test_dir):
        try:
            os.makedirs(test_dir)
        except OSError as oe:
            if oe.errno != errno.EEXIST:
                raise


def get_tags(username, repo, registry=None, auth=None):
    if registry != "registry.hub.docker.com":
        if "://" not in registry:
            registry = "https://%s" % registry
        auth = DOCKER_CREDS.get(registry, {})
        hub = dockerhub.DockerHub(url=registry, namespace=username, repo=repo, **auth)
    else:
        auth = DOCKER_CREDS.get('https://registry.hub.docker.com', {})
        hub = dockerhub.DockerHub(**auth)
    tag_names = []
    for tag in hub.tags(username, repo):
        tag_names.append(tag['name'])
    return tag_names

def publish_txn():
    global _in_txn
    if _in_txn:
        _in_txn = False
        return os.system("cvmfs_server publish singularity.opensciencegrid.org")
    return 0


def make_final_symlink(image_dir, singularity_rootfs, namespace, repo_name, repo_tag):
    """
    Create symlink: $ROOTFS/.images/$HASH -> $ROOTFS/$NAMESPACE/$IMAGE:$TAG
    """
    final_path = os.path.join(singularity_rootfs, namespace, "%s:%s" % (repo_name, repo_tag))
    final_dir = os.path.split(final_path)[0]
    if not os.path.exists(final_dir):
        retval = start_txn(singularity_rootfs)
        if retval:
            return retval
        try:
            os.makedirs(final_dir)
        except OSError as oe:
            if oe.errno != errno.EEXIST:
                raise
    if os.path.exists(final_path):
        # Final symlink exists and is already correct.
        link_value = os.readlink(final_path)
        if link_value == image_dir:
            print("Image is already latest revision.")
            return 0
        # Otherwise, clear the symlink; we will recreate.  Since CVMFS is transactional,
        # we don't care that an unlink / symlink is not atomic.
        retval = start_txn(singularity_rootfs)
        if retval:
            return retval
        os.unlink(final_path)
    retval = start_txn(singularity_rootfs)
    if retval:
        return retval
    os.symlink(image_dir, final_path)
    return 0


def parse_image(image):
    """
    Parse an image string into image components, setting appropriate defaults.
    Returns a tuple of (registry, namespace, repo_name, repo_tag).
    """

    # INPUT PARSING -------------------------------------------
    # Parse registry, image name, repo name, and namespace
    # Support different formats:
    # opensciencegrid/osgvo-julia:latest
    # opensciencegrid/osgvo-ants
    # openjdk:8
    # containers.ligo.org/lscsoft/lalsuite/lalsuite-v6.53:stretch
    registry = "registry.hub.docker.com"
    # First split the docker image name by /
    split_image = image.split('/')

    # If there are two parts, we have namespace with repo (and maybe tab)
    if len(split_image) == 2:
        namespace = split_image[0]
        image = split_image[1]
    elif len(split_image) > 2:
        # We have a custom registry
        registry = split_image[0]
        #print("Custom registry:", registry)
        namespace = split_image[1]
        image = "/".join(split_image[2:])
    # Otherwise, we must be using library namespace
    else:
        namespace = "library"
        image = split_image[0]

    # Now split the docker image name by :
    image = image.split(':')
    if len(image) == 2:
        repo_name = image[0]
        repo_tag = image[1]

    # Otherwise, assume latest of an image
    else:
        repo_name = image[0]
        repo_tag = "latest"

    return registry, namespace, repo_name, repo_tag

def get_manifest(hub, namespace, repo_name, repo_tag, manifest_cache):
    metadata = hub.manifest(namespace, repo_name, repo_tag, head=True)
    digest = metadata.headers['docker-content-digest']

    if digest in manifest_cache:
        return manifest_cache[digest], digest
    else:
        manifest = hub.manifest(namespace, repo_name, repo_tag)
        manifest_cache[digest] = manifest
        return manifest, digest

def publish_image(image, singularity_rootfs, registry, doauth, manifest_cache):

    # Tell the user the namespace, repo name and tag
    registry, namespace, repo_name, repo_tag = parse_image(image)
    print("Docker image to publish: %s/%s/%s:%s" % (registry, namespace, repo_name, repo_tag))


# IMAGE METADATA -------------------------------------------
# Use Docker Registry API (version 2.0) to get images ids, manifest

    # Get an image manifest - has image ids to parse, and will be
    # used later to get Cmd
    # Prepend "https://" to the registry
    if "://" not in registry:
        registry = "https://%s" % registry
    auth = DOCKER_CREDS.get(registry, {})
    hub = dockerhub.DockerHub(url=registry, namespace=namespace, repo=repo_name, **auth)
    manifest, digest = get_manifest(hub, namespace, repo_name, repo_tag, manifest_cache)

    # Calculate a unique hash across all layers.  We'll use that as the identifier
    # for the final image.
    if manifest['schemaVersion'] == 1 and 'fsLayers' in manifest:
        hasher = hashlib.sha256()
        for layerinfo in manifest['fsLayers']:
            layer_hash = layerinfo['blobSum'].split(":")[-1].encode('utf-8')
            hasher.update(layer_hash)
        image_hash = hasher.hexdigest()
    elif manifest['schemaVersion'] == 2 and 'layers' in manifest:
        hasher = hashlib.sha256()
        for layerinfo in manifest['layers']:
            layer_hash = layerinfo['digest'].split(":")[-1].encode('utf-8')
            hasher.update(layer_hash)
        image_hash = hasher.hexdigest()
    else:
        # Use the Docker-Content-Digest to identify the image
        image_hash = digest

    # The image is staged to $ROOTFS/.images/$HASH
    # Docker-Content-Digest uses format 'sha256:dead...beef'
    # Split it after the hash type       ^^^^^^^^^
    #   plus two chars                            ^^^^^^^^^
    sep = image_hash.find(':')
    image_parentdir = os.path.join(singularity_rootfs, ".images", image_hash[0:sep+3])
    image_dir = os.path.join(image_parentdir, image_hash[sep+3:])

    # If the image has already been staged, simply return.
    if os.path.exists(image_dir):
        make_final_symlink(image_dir, singularity_rootfs, namespace, repo_name, repo_tag)
        return publish_txn()
    else:
        print("Image dir, %s, does not exist; triggering CVMFS mount." % image_dir)
        retval = start_txn(singularity_rootfs)
        if os.path.exists(image_dir):   # Same as above
            make_final_symlink(image_dir, singularity_rootfs, namespace, repo_name, repo_tag)
            return publish_txn()
        if retval:
            return retval
        try:
            # Create the parent directory, but not the image directory itself
            # Singularity will create the image directory on successful conversion
            os.makedirs(image_parentdir)
        except OSError as oe:
            if oe.errno != errno.EEXIST:
                raise

#  DOWNLOAD IMAGE --------------------------------------------
    # Use /var/tmp for singularity temporary files
    singularity_env = os.environ.copy()
    singularity_env['TMPDIR'] = '/var/tmp'

    if 'username' in auth:
        singularity_env['SINGULARITY_DOCKER_USERNAME'] = auth['username']
    if 'password' in auth:
        singularity_env['SINGULARITY_DOCKER_PASSWORD'] = auth['password']

    print("Calling singularity to build sandbox from image")
    subprocess.check_call(
        ['singularity', '--silent', 'build',
         '--disable-cache=true', # Images are only downloaded once
         '--force',              # Don't get stuck at a prompt if the target somehow exists
         '--fix-perms',
         '--sandbox',
         image_dir, 'docker://' + image],
        env=singularity_env)

    # Various fixups to make the image compatible with CVMFS and singularity.
    srv = os.path.join(image_dir, "srv")
    cvmfs = os.path.join(image_dir, "cvmfs")
    if not os.path.exists(srv):
        os.makedirs(srv)
    if not os.path.exists(cvmfs):
        os.makedirs(cvmfs)

    make_final_symlink(image_dir, singularity_rootfs, namespace, repo_name, repo_tag)
    # Publish CVMFS as necessary.
    return publish_txn()

def verify_image(image, registry, doauth, manifest_cache):

    # Tell the user the namespace, repo name and tag
    registry, namespace, repo_name, repo_tag = parse_image(image)
    print("Docker image to verify: %s/%s/%s:%s" % (registry, namespace, repo_name, repo_tag))

# IMAGE METADATA -------------------------------------------
# Use Docker Registry API (version 2.0) to get images ids, manifest

    # Get an image manifest - has image ids to parse, and will be
    # used later to get Cmd
    # Prepend "https://" to the registry
    if "://" not in registry:
        registry = "https://%s" % registry
    auth = DOCKER_CREDS.get(registry, {})
    hub = dockerhub.DockerHub(url=registry, namespace=namespace, repo=repo_name, **auth)
    retval = 0
    try:
        hub.manifest(namespace, repo_name, repo_tag, head=True)
        print(repo_name + ":" + repo_tag + " manifest found")
        retval = 0
    except Exception as ex:
        print(repo_name + ":" + repo_tag + " manifest not found")
        traceback.print_exc()
        retval = 1
        raise
    return retval

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