#!/usr/bin/env python

# Copyright 2017 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.

import argparse
import json
import os
import shutil
import struct
import subprocess
import sys

def gen_pkg_key(pm_bin, dest):
    """Create a key to use for signing FARs.

    pm_bin: path to the 'pm' binary
    dest  : path where the key should be written

    If successful None is returned, otherwise a string describing the error
    encountered is returned.
    """
    pm_cmd = [pm_bin, "-k", dest, "genkey"]
    try:
        subprocess.check_call(pm_cmd)
        return None
    except subprocess.CalledProcessError as e:
        return "Error generating FAR signing key: %s" % e
    except OSError as e:
        return "Error launching PM binary: %s" % e

def build_package(pm_bin, pkg_key, far_stg_dir, manifest, pkg_name):
    """Build a metadata FAR describing the package.

    pm_bin     : path to the 'pm' binary
    pkg_key    : path to the key to use to sign the FAR
    far_stg_dir: a working directory to use for staging files before they are
                 written into a single file
    manifest   : path to a manifest file describing the contents of the package
    pkg_name   : the name of the package

    Create a signed metadata FAR representing the package. This function
    returns a tuple which is the path to the metadata FAR and any error. If
    there is an error the first member of the tuple is None and the second
    member contains a string describing the error. If there is no error, the
    second member of tuple is None.
    """
    init_cmd = [pm_bin, "-o", far_stg_dir, "-n", pkg_name, "init"]
    try:
        subprocess.check_call(init_cmd)
    except subprocess.CalledProcessError as e:
        return None, None, "Could not initialize package: %s" % e
    except OSError as e:
        return None, None, "Could not start package initializer: %s" % e

    build_cmd = [pm_bin, "-o", far_stg_dir, "-k", pkg_key, "-m", manifest, "build"]
    try:
        subprocess.check_call(build_cmd)
    except subprocess.CalledProcessError as e:
        return None, "Could not create package metadata FAR: %s" % e
    except OSError as e:
        return None, None, "Could not start packging tool %s" % e

    far_path = os.path.join(far_stg_dir, "meta.far")
    pkg_json = os.path.join(far_stg_dir, "meta", "package")
    if os.path.exists(far_path) and os.path.exists(pkg_json):
        return far_path, pkg_json, None
    else:
        return None, None, "Unknown failure, metadata package not produced"

def assemble_manifest(manifests_dir, output_stream):
    """Create a single manifest from the joining of the system and boot
    manifests.

    manifests_dir: Directory which contains system_manifest and/or
                   boot_manifest.
    output_stream: An output stream where the combined manifest will be
                   written.

    No returns, but may raise an exception if writing to output_stream fails.
    """
    inputs = ["final_package_manifest"]
    for input in inputs:
        manifest = os.path.join(manifests_dir, input)
        if os.path.exists(manifest) and os.stat(manifest).st_size > 0:
            with open(manifest, "r") as src:
                for line in src:
                    output_stream.write(line)
            break

def add_far_to_repo(amber_bin, name, far, key_dir, repo_dir, version=0):
    """Add a FAR to the update repository under the specified name

    amber_bin: path to the amber binary
    name     : name to publish the package as
    far      : path to the FAR file to publish
    key_dir  : directory containing set of keys for the update repository
    repo_dir : directory to use as the update repository (should exist, but
               doesn't need to be initialized as an update repo)

    On success returns None, otherwise returns a string describing the error
    that occurred.
    """
    cmd = [amber_bin, "-r", repo_dir, "-p", "-f", far, "-n", "%s/%d" % (name, version), "-k", key_dir]

    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        return "Failure running package publisher: %s" % e
    except OSError as e:
        return "Failure launching package publisher: %s" % e

    return None

def add_rsrcs_to_repo(amber_bin, manifest, key_dir, repo_dir):
    """Add package resources, aka. content blobs to the update repository.

    amber_bin: path to the amber binary
    manifest : a file containing a mapping of file paths on the target to
               paths on the host. All paths will be added to the update
               respository, named after their content ID.
    key_dir  : directory containing set of keys for the update repository
    repo_dir : directory to use as the update repository (should exist, but
               doesn't need to be initialized as an update repo)

    On success returns None, otherwise returns a string describing the error
    that occurred.
    """
    cmd = [amber_bin, "-r", repo_dir, "-m", "-f", manifest, "-k", key_dir]

    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        return "Failure adding content blob: %s" % e
    except OSError as e:
        return "Unable to launch blob tool: %s" % e

    return None

def publish(pm_bin, amber_bin, pkg_key, repo_key_dir, pkg_stg_dir, update_repo,
            manifests_dir, pkgs, verbose):
    """Publish packages as a signed metadata FAR and a collection of content
    blobs named after their content IDs.

    pm_bin       : path to the pm binary
    amber_bin    : path to the amber binary
    pkg_key      : path to the key to use to sign the metdata FARs
    repo_key_dir : directory containing keys to use for the update respository
    pkg_stg_dir  : a directory that can be used for staging temporary files
                   when creating the packages
    update_repo  : a path that is or can be used as the update repository
    manifests_dir: parent directory where manifests can be found for the
                   packages
    pkgs         : pkgs to publish
    verbose      : print out status updates
    """
    # as an optimization, publish the content blobs all at once
    master_fest = os.path.join(pkg_stg_dir, "masterfest")
    master_fd = open(master_fest, "w+")
    count = len(pkgs)
    for pkg in pkgs:
        # this process can be time consuming with a large number of packages
        # so printing incremental progress can be reassuring
        if verbose:
            sys.stdout.write("Packages remaining: %d       \r" % count)
            sys.stdout.flush()

        far_base = os.path.join(pkg_stg_dir, pkg)
        far_stg = os.path.join(far_base, "archive")

        if os.path.exists(far_base):
            shutil.rmtree(far_base)
        os.makedirs(far_stg)

        manifest = os.path.join(far_base, "manifest")
        src_manifest_dir = os.path.join(manifests_dir, pkg)
        with open(manifest, "w+") as man_fd:
            assemble_manifest(src_manifest_dir, man_fd)

        # some packages are apparently devoid of content, skip them
        if os.stat(manifest).st_size == 0:
            count -= 1
            continue

        with open(manifest, "r") as man_fd:
            for line in man_fd:
                master_fd.write(line)

        meta_far, pkg_json, err = build_package(pm_bin, pkg_key, far_stg, manifest, pkg)
        if err is not None:
            print "Building package failed: %s" % err
            break

        pkg_version = None
        with open(pkg_json, 'r') as pkg_meta:
            meta = json.load(pkg_meta)
            pkg_version = int(meta["version"])

        if pkg_version is None:
            print "Could not read version from %q" % pkg_json
            break

        result = add_far_to_repo(amber_bin, pkg, meta_far, repo_key_dir, update_repo,
                                 version=pkg_version)
        if result is not None:
            print "Package not added to update repo: %s" % result
            break
        count -= 1

    master_fd.close()
    result = add_rsrcs_to_repo(amber_bin, master_fest, repo_key_dir, update_repo)
    if result is not None:
        print "Package contents not added to update repo: %s" % result
        return -1

    return count

def main():
    parser = argparse.ArgumentParser(description=("Publish one or more build "
                                                  "packages as package manager "
                                                  "packages."))
    parser.add_argument('--build-dir', action='store', required=True)
    parser.add_argument('--update-repo', action='store', required=False)
    parser.add_argument('--update-keys', action='store', required=False)
    parser.add_argument('--pkg-key', action='store', required=False)
    parser.add_argument('--fars-dir', action='store', required=False,
                        help="""Directory where intermediate files for the
                        package(s) will be stored""")
    parser.add_argument('--pkgs', action='append', required=False,
                        help="""Packages to publish. This argument may be
                        repeated to publish multiple packages.""")
    parser.add_argument('--quiet', action='store_true', required=False, default=False)
    args = parser.parse_args()

    build_dir = args.build_dir
    ptr_size = 8 * struct.calcsize("P")
    host_tools_dir = os.path.join(build_dir, "host_x%d" % ptr_size)
    pm_bin = os.path.join(host_tools_dir, "pm")
    if not os.path.exists(pm_bin):
        print "Could not find 'pm' tool at %s" % pm_bin
        return -1

    amber_bin = os.path.join(host_tools_dir, "amber-publish")
    if not os.path.exists(amber_bin):
        print "Could not find amber-publish tool at %s" % amber_bin
        return -1

    repo_dir = args.update_repo
    if not repo_dir:
        repo_dir = os.path.join(build_dir, "amber-files")
        if not os.path.exists(repo_dir):
          os.makedirs(repo_dir)

    if not os.path.exists(repo_dir):
        print "Publishing repository directory '%s' could not be found" % repo_dir
        return -1

    keys_src_dir = args.update_keys
    if not keys_src_dir:
        keys_src_dir = build_dir

    pkg_key = args.pkg_key
    if not pkg_key:
        pkg_key = os.path.join(build_dir, "pkg_key")
        result = gen_pkg_key(pm_bin, pkg_key)
        if result is not None:
            print result
            return -1

    pkg_stg_dir = args.fars_dir
    if not pkg_stg_dir:
        pkg_stg_dir = os.path.join(build_dir, "fars")
        if not os.path.exists(pkg_stg_dir):
          os.makedirs(pkg_stg_dir)

    if not os.path.exists(pkg_stg_dir):
        print "Packages staging directory '%s' could not be found" % pkg_stg_dir
        return -1

    pkg_list = args.pkgs
    if not pkg_list:
        pkg_list = []
        list_path = os.path.join(build_dir, "gen", "build", "gn", "packages")
        with open(list_path, "r") as pfile:
            for l in pfile:
                pkg_list.append(l.strip())

    return publish(pm_bin, amber_bin, pkg_key, build_dir, pkg_stg_dir, repo_dir,
        os.path.join(build_dir, "package"), pkg_list, not args.quiet)

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