iat-verifier/scripts/compile_token

#!/usr/bin/env python3
#-------------------------------------------------------------------------------
# Copyright (c) 2019-2022, Arm Limited. All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#
#-------------------------------------------------------------------------------

import argparse
import logging
import os
import sys

from ecdsa import SigningKey
from iatverifier.util import read_token_map, convert_map_to_token
from iatverifier.psa_iot_profile1_token_verifier import PSAIoTProfile1TokenVerifier
from iatverifier.verifiers import AttestationTokenVerifier


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument('source', help='Token source in YAML format')
    parser.add_argument('-o', '--outfile',
                        help='''Output file for the compiled token. If this is not
                        specified, the token will be written to standard output.''')
    parser.add_argument('-k', '--keyfile',
                        help='''Path to the key in PEM format that should be used to
                        sign the token. If this is not specified, the token will be
                        unsigned.''')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-r', '--raw', action='store_true',
                       help='''Generate raw CBOR and do not create a signature
                       or COSE wrapper.''')
    group.add_argument('-m', '--hmac', action='store_true',
                       help='''Generate a token wrapped in a Mac0 rather than
                       Sign1 COSE structure.''')

    args = parser.parse_args()
    signing_key = None

    cose_alg = None
    if args.hmac:
        method = AttestationTokenVerifier.SIGN_METHOD_MAC0
        cose_alg = AttestationTokenVerifier.COSE_ALG_HS256

        if args.keyfile:
            with open(args.keyfile, 'rb') as fh:
                signing_key = fh.read()
    elif args.raw:
        if args.keyfile:
            raise ValueError('A keyfile cannot be specified with --raw.')
        method = AttestationTokenVerifier.SIGN_METHOD_RAW
    else:
        method = AttestationTokenVerifier.SIGN_METHOD_SIGN1
        if args.keyfile:
            with open(args.keyfile) as fh:
                signing_key = SigningKey.from_pem(fh.read())

    verifier = PSAIoTProfile1TokenVerifier.get_verifier()
    if verifier.method != method:
        verifier.method = method
    if cose_alg is not None and verifier.cose_alg != cose_alg:
        verifier.cose_alg = cose_alg
    token_map = read_token_map(args.source)

    if args.outfile:
        with open(args.outfile, 'wb') as wfh:
            convert_map_to_token(token_map, signing_key, verifier, wfh)
    else:
        with os.fdopen(sys.stdout.fileno(), 'wb') as wfh:
            convert_map_to_token(token_map, signing_key, verifier, wfh)