bazel_rules_fuchsia/tools/artifact/fetch.py

#!/usr/bin/env python3
# Copyright 2021 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Fetches a package.far as specified in the artifact_lock.json file."""

import argparse
import json
import os
import subprocess
import shutil
import sys
import tempfile
import urllib.error
import urllib.request

META_FAR_FILE_RELATIVE_PATH = 'meta.far'
CONTENTS_FILE_RELATIVE_PATH = 'meta/contents'

BLOB_SERVER = 'fuchsia-blobs.googleusercontent.com'


def extract_far(far_tool, meta_far_path, extract_dir):
    """Invokes the far tool to extract the meta_far."""
    args = [
        far_tool, 'extract', '--archive=' + meta_far_path,
        '--output=' + extract_dir
    ]
    subprocess.check_output(args, stderr=subprocess.STDOUT)


def get_merkle(merkleroot_tool, path):
    """Run the merkleroot tool to get the merkle."""
    result = subprocess.run(
        [merkleroot_tool, path],
        check=True,
        capture_output=True,
        text=True,
    )
    return result.stdout.split()[0]


def compare_merkle(merkleroot_tool, meta_far_path, meta_far_merkle):
    """Compares the saved merkle to see if download is needed."""
    if not os.path.isfile(meta_far_path):
        return False
    return meta_far_merkle == get_merkle(merkleroot_tool, meta_far_path)


def parse_args():
    """Parses arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument('--far-tool', help='Path to far tool', required=True)
    parser.add_argument(
        '--merkleroot-tool', help='Path to merkleroot tool', required=True)
    parser.add_argument(
        '--lock-file', help='Path to saved lock file', required=True)
    parser.add_argument(
        '--outdir', help='Path to target output directory', required=True)
    parser.add_argument(
        '--local-dir', help='Path to local artifact_store root', required=False)
    parser.add_argument(
        '--show-progress', dest='progress', action='store_true',
        help='Show progress during execution',
        required=False)
    return parser.parse_args()


def fetch_file(artifact_store, merkle, relative_path, name, args):
    """Fetch the merkle contents from the artifact store.

    In the case of local store the output relative path and name
    are used for input since the in/out directory structures are
    identical.
"""

    # Make sure the output path exists
    path = os.path.join(args.outdir, relative_path)
    os.makedirs(path, exist_ok=True)
    filename = os.path.join(path, name)

    if args.progress:
        print(" - %s" % filename)

    type_ = artifact_store['type']
    if type_ == 'tuf':
        url = 'https://%s/%s' % (BLOB_SERVER, merkle)
        try:
            urllib.request.urlretrieve(url, filename)
        except urllib.error.HTTPError as ex:
            raise ValueError('Failed to download meta far or blob') from ex
    elif type_ == 'local':
        local_filename = os.path.join(args.local_dir, relative_path, name)
        shutil.copyfile(local_filename, filename)
    else:
        raise TypeError('Unknown artifact store type: ' + type_)


def fetch_artifact(artifact, args):
    """Fetch the meta-far and blobs."""
    merkle = artifact['id']
    artifact_store = artifact['artifact_store']
    artifact_name = artifact['name'].replace('/0', '')
    artifact_type = artifact['type']

    # meta.far is saved under artifact_stores/<artifact_store_name>/<artifact_name>
    relative_path = os.path.join('artifact_stores', artifact_store['name'])
    meta_far_filename = os.path.join(args.outdir, relative_path, artifact_name)

    if compare_merkle(args.merkleroot_tool, meta_far_filename, merkle):
        return

    fetch_file(artifact_store, merkle, relative_path, artifact_name, args)

    if artifact_type == 'package':
        with tempfile.TemporaryDirectory() as tmpdir:
            extract_far(args.far_tool, meta_far_filename, tmpdir)
            with open(os.path.join(tmpdir, CONTENTS_FILE_RELATIVE_PATH),
                      'r') as content_file:
                for content in content_file.readlines():
                    merkle = content.strip().split('=')[1]
                    fetch_file(artifact_store, merkle, 'blobs', merkle, args)


def main():
    """Main for the fetch tool."""
    args = parse_args()
    os.makedirs(args.outdir, exist_ok=True)

    with open(args.lock_file, 'r') as lock_file:
        try:
            artifact_lock = json.load(lock_file)
        except ValueError as ex:
            raise ValueError(
                'Malformed %s: %s\n%s\n' %
                (args.lock_file, lock_file.read(),
                 'Execute artifact_update.py to refresh merkles.')) from ex
        for artifact in artifact_lock['artifacts']:
            fetch_artifact(artifact, args)


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