# Copyright 2021 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Metadata signing facilities."""

import argparse
from pathlib import Path

from pw_software_update import keys
from pw_software_update.tuf_pb2 import SignedRootMetadata
from pw_software_update.update_bundle_pb2 import UpdateBundle


def sign_root_metadata(
    root_metadata: SignedRootMetadata, root_key_pem: bytes
) -> SignedRootMetadata:
    """Signs or re-signs a Root Metadata.

    Args:
      root_metadata: A SignedRootMetadata to be signed/re-signed.
      root_key_pem: The Root signing key in PEM.
    """

    signature = keys.create_ecdsa_signature(
        root_metadata.serialized_root_metadata, root_key_pem
    )
    root_metadata.signatures.append(signature)

    return root_metadata


def sign_update_bundle(
    bundle: UpdateBundle, targets_key_pem: bytes
) -> UpdateBundle:
    """Signs or re-signs an update bundle.

    Args:
      bundle: An UpdateBundle to be signed/re-signed.
      targets_key_pem: The targets signing key in PEM.
    """
    bundle.targets_metadata['targets'].signatures.append(
        keys.create_ecdsa_signature(
            bundle.targets_metadata['targets'].serialized_targets_metadata,
            targets_key_pem,
        )
    )
    return bundle


def parse_args():
    """Parse CLI arguments."""
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument(
        '--root-metadata',
        type=Path,
        required=False,
        help='Path to the root metadata to be signed',
    )

    parser.add_argument(
        '--bundle',
        type=Path,
        required=False,
        help='Path to the bundle to be signed',
    )

    parser.add_argument(
        '--output',
        type=Path,
        required=False,
        help=(
            'Path to save the signed root metadata or bundle '
            'to; Defaults to the input path if unspecified'
        ),
    )

    parser.add_argument(
        '--key', type=Path, required=True, help='Path to the signing key'
    )

    args = parser.parse_args()

    if not (args.root_metadata or args.bundle):
        parser.error('either "--root-metadata" or "--bundle" must be specified')
    if args.root_metadata and args.bundle:
        parser.error('"--root-metadata" and "--bundle" are mutually exclusive')

    return args


def main(root_metadata: Path, bundle: Path, key: Path, output: Path) -> None:
    """Signs or re-signs a root metadata or an update bundle."""
    if root_metadata:
        signed_root_metadata = sign_root_metadata(
            SignedRootMetadata.FromString(root_metadata.read_bytes()),
            key.read_bytes(),
        )

        if not output:
            output = root_metadata
        output.write_bytes(signed_root_metadata.SerializeToString())
    else:
        signed_bundle = sign_update_bundle(
            UpdateBundle.FromString(bundle.read_bytes()), key.read_bytes()
        )

        if not output:
            output = bundle
        output.write_bytes(signed_bundle.SerializeToString())


if __name__ == '__main__':
    main(**vars(parse_args()))