iat-verifier/scripts/compile_token

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

"""CLI tool for compiling token from a yaml file"""

import argparse
import logging
import os
import sys

from iatverifier.util import read_token_map, convert_map_to_token, read_keyfile
from iatverifier.util import get_cose_alg_from_key
from iatverifier.psa_iot_profile1_token_verifier import PSAIoTProfile1TokenVerifier
from iatverifier.psa_2_0_0_token_verifier import PSA_2_0_0_TokenVerifier
from iatverifier.attest_token_verifier import AttestationTokenVerifier, VerifierConfiguration
from iatverifier.cca_token_verifier import CCATokenVerifier, CCAPlatformTokenVerifier

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

    token_verifiers = {
        "PSA-IoT-Profile1-token": PSAIoTProfile1TokenVerifier,
        "CCA-token": CCATokenVerifier,
        "CCA-plat-token": CCAPlatformTokenVerifier,
        "PSA-2.0.0-token": PSA_2_0_0_TokenVerifier,
    }

    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('--psa-iot-profile1-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.''')
    parser.add_argument('--cca-platform-token-keyfile',
                        help='''Path to the key in PEM format that should be used to
                        sign the CCA platform token. If this is not specified,
                        the token will be unsigned.''')
    parser.add_argument('--cca-realm-token-keyfile',
                        help='''Path to the key in PEM format that should be used to
                        sign the CCA Realm token. If this is not specified, the
                        token will be unsigned.''')
    group = parser.add_mutually_exclusive_group()
    parser.add_argument('-a', '--add-protected-header', action='store_true',
                        help='''
                        Add protected header to the COSE wrapper.
                        ''')
    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.''')
    parser.add_argument('-t', '--token-type',
                        help='''The type of the Token.''',
                        choices=token_verifiers.keys(),
                        required=True)

    args = parser.parse_args()

    if args.hmac:
        METHOD = AttestationTokenVerifier.SIGN_METHOD_MAC0
    elif args.raw:
        if args.psa_iot_profile1_keyfile:
            raise ValueError('A keyfile cannot be specified with --raw.')
        METHOD = AttestationTokenVerifier.SIGN_METHOD_RAW
    else:
        METHOD = AttestationTokenVerifier.SIGN_METHOD_SIGN1

    configuration = VerifierConfiguration(strict=True, keep_going=False)

    verifier_class = token_verifiers[args.token_type]
    if verifier_class == PSAIoTProfile1TokenVerifier:
        key = read_keyfile(args.psa_iot_profile1_keyfile, METHOD)
        if METHOD == AttestationTokenVerifier.SIGN_METHOD_SIGN1:
            cose_alg = get_cose_alg_from_key(
                key,
                AttestationTokenVerifier.COSE_ALG_ES256)
        else:
            cose_alg = AttestationTokenVerifier.COSE_ALG_HS256
        verifier = PSAIoTProfile1TokenVerifier(
            method=METHOD,
            cose_alg=cose_alg,
            signing_key=key,
            configuration=configuration)
    elif verifier_class == CCATokenVerifier:
        if METHOD != AttestationTokenVerifier.SIGN_METHOD_SIGN1:
            logging.error('Only sign1 method is supported by this token type.\n\t')
            sys.exit(1)
        platform_token_key = read_keyfile(args.cca_platform_token_keyfile, METHOD)
        realm_token_key = read_keyfile(args.cca_realm_token_keyfile, METHOD)
        realm_token_method = AttestationTokenVerifier.SIGN_METHOD_SIGN1
        platform_token_method = AttestationTokenVerifier.SIGN_METHOD_SIGN1
        realm_token_cose_alg = get_cose_alg_from_key(
            realm_token_key,
            AttestationTokenVerifier.COSE_ALG_ES384)
        platform_token_cose_alg = get_cose_alg_from_key(
            platform_token_key,
            AttestationTokenVerifier.COSE_ALG_ES384)
        verifier = CCATokenVerifier(
            realm_token_method=realm_token_method,
            realm_token_cose_alg=realm_token_cose_alg,
            realm_token_key=realm_token_key,
            platform_token_method=platform_token_method,
            platform_token_cose_alg=platform_token_cose_alg,
            platform_token_key=platform_token_key,
            configuration=configuration)
    elif verifier_class == CCAPlatformTokenVerifier:
        key_checked = args.cca_platform_token_keyfile
        key = read_keyfile(args.cca_platform_token_keyfile, METHOD)
        cose_alg = get_cose_alg_from_key(key, AttestationTokenVerifier.COSE_ALG_ES384)
        verifier = CCAPlatformTokenVerifier(
            method=AttestationTokenVerifier.SIGN_METHOD_SIGN1,
            cose_alg=cose_alg,
            signing_key=key,
            configuration=configuration,
            necessity=None)
    elif verifier_class == PSA_2_0_0_TokenVerifier:
        key_checked = args.psa_iot_profile1_keyfile
        key = read_keyfile(keyfile=args.psa_iot_profile1_keyfile, method=METHOD)
        if METHOD == AttestationTokenVerifier.SIGN_METHOD_SIGN1:
            cose_alg = get_cose_alg_from_key(key, AttestationTokenVerifier.COSE_ALG_ES256)
        else:
            cose_alg = AttestationTokenVerifier.COSE_ALG_HS256
        verifier = PSA_2_0_0_TokenVerifier(
            method=METHOD,
            cose_alg=cose_alg,
            signing_key=key,
            configuration=configuration)
    else:
        logging.error(f'Invalid token type:{verifier_class}\n\t')
        sys.exit(1)
    token_map = read_token_map(args.source)

    if args.outfile:
        with open(args.outfile, 'wb') as wfh:
            convert_map_to_token(
                token_map,
                verifier,
                wfh,
                add_p_header=args.add_protected_header,
                name_as_key=True,
                parse_raw_value=True)
    else:
        with os.fdopen(sys.stdout.fileno(), 'wb') as wfh:
            convert_map_to_token(
                token_map,
                verifier,
                wfh,
                add_p_header=args.add_protected_header,
                name_as_key=True,
                parse_raw_value=True)