image-upload

#!/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/>.

import argparse
import os
import re
import subprocess
import sys
import time
import urllib.parse

from lib import s3
from lib.constants import BOTS_DIR, IMAGES_DIR
from lib.directories import get_images_data_dir
from lib.network import get_curl_ca_arg, redhat_network
from lib.stores import PUBLIC_STORES, REDHAT_STORES


def upload(store: str, source: str, public: bool, prune: bool = False) -> bool:
    dest = urllib.parse.urljoin(store, os.path.basename(source))
    url = urllib.parse.urlparse(dest)

    if not s3.is_key_present(url):
        # No credentials?  That's not going to work...
        sys.stderr.write(f"image-upload: no credentials for {dest}")
        return False

    # Start building the command
    cmd = ["curl", "--no-progress-meter", "--fail", "--upload-file", source, *get_curl_ca_arg(url.netloc)]

    headers = {s3.ACL: s3.PUBLIC} if public else {}

    # slightly magic: we already know the checksum: it's part of the filename
    match = re.search('[0-9a-f]{64}', source)
    assert match is not None
    hash_value = match.group(0)

    cmd += s3.sign_curl(url, method='PUT', headers=headers, checksum=hash_value)

    # Only prune after the above checks, to ensure that we have credentials.
    if prune:
        subprocess.check_call([f'{BOTS_DIR}/image-prune', '--s3', store])

    print("Uploading to", dest, file=sys.stderr)

    # Retry with exponential back-off, to defend against short-term networking errors
    for retry in range(3):
        try:
            subprocess.check_call(cmd)
            return True
        except subprocess.CalledProcessError:
            delay = 10 * 4 ** retry
            sys.stderr.write(f"image-upload: retrying upload to {dest} in {delay}s..\n")
            time.sleep(delay)

    sys.stderr.write(f"image-upload: unable to upload image: {dest}\n")
    return False


def main() -> None:
    parser = argparse.ArgumentParser(description='Upload bot state or images')
    parser.add_argument("--prune-s3", action="store_true", help="Run image-prune before upload to S3")
    parser.add_argument("--store", action="append", default=[], help="Where to send state or images")
    parser.add_argument("--state", action="store_true", help="Images or state not recorded in git")
    parser.add_argument('image', nargs='*')
    args = parser.parse_args()

    data_dir = get_images_data_dir()
    sources = []
    for image in args.image:
        if args.state:
            source = os.path.join(data_dir, image)
        else:
            link = os.path.join(IMAGES_DIR, image)
            if not os.path.islink(link):
                parser.error("image link does not exist: " + image)
            source = os.path.join(data_dir, os.readlink(link))
        if not os.path.isfile(source):
            parser.error("image does not exist: " + image)
        sources.append(source)

    for source in sources:
        public = not os.path.basename(source).startswith('rhel')

        stores = args.store
        if not stores:
            if image_upload_store := os.environ.get('COCKPIT_IMAGE_UPLOAD_STORE'):
                # NB: This is a *complete replacement* of the built-in stores,
                # used for testing.  Compare with image-download which only
                # adds this to the list, if it's present.
                stores = [image_upload_store]
        if not stores:
            stores = list(PUBLIC_STORES)
            if redhat_network():
                stores += REDHAT_STORES
        success = False
        for store in stores:
            success |= upload(store, source, public, prune=args.prune_s3)

        if not success:
            sys.exit('Failed to upload to any image store')


if __name__ == '__main__':
    main()