#!/usr/bin/env python3
"""
Generate test cases for 'pk_can_do_psa()' in order to exercize
'mbedtls_pk_can_do_psa()'.

Generated test file is saved in 'tests/suites/test_suite_pk.pk_can_do_psa.data'
for the time being.
"""

## Copyright The Mbed TLS Contributors
## SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later

import os
import re

HEADER_MESSAGE = f"""\
# This files was automatically generated by {os.path.basename(__file__)}.
# DO NOT edit it manually.\n
"""

# PK contexts to be tested
# pylint: disable=line-too-long
PK_CONTEXTS = {
    'MBEDTLS_PK_ECDSA': {
        'PSA_KEY_TYPE_ECC_KEY_PAIR(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
        },
        'PSA_KEY_TYPE_ECC_PUBLIC_KEY(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            # Skip context with PSA_ALG_SHA_256 because pk_can_do_psa() cannot verify which alg + hash_alg
            # is associated with public key.
        },
    },
    'MBEDTLS_PK_ECKEY': {
        'PSA_KEY_TYPE_ECC_KEY_PAIR(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE',
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
        'PSA_KEY_TYPE_ECC_PUBLIC_KEY(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            # Skip context with PSA_ALG_SHA_256 because pk_can_do_psa() cannot verify which alg + hash_alg
            # is associated with public key.

            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
    },
    'MBEDTLS_PK_ECKEY_DH': {
        'PSA_KEY_TYPE_ECC_KEY_PAIR(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE',
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
        'PSA_KEY_TYPE_ECC_PUBLIC_KEY(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
    },
    'MBEDTLS_PK_RSA': {
        'PSA_KEY_TYPE_RSA_KEY_PAIR:1024': {
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_CRYPT': [
                'PSA_KEY_USAGE_DECRYPT',
                'PSA_KEY_USAGE_ENCRYPT',
            ],
            'PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_DECRYPT',
                'PSA_KEY_USAGE_ENCRYPT',
            ],
        },
        'PSA_KEY_TYPE_RSA_PUBLIC_KEY:1024': {
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_CRYPT': [
                'PSA_KEY_USAGE_ENCRYPT',
            ],
            'PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_ENCRYPT',
            ],
        },
    },
    'MBEDTLS_PK_OPAQUE': {
        'PSA_KEY_TYPE_ECC_KEY_PAIR(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE',
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
        'PSA_KEY_TYPE_ECC_PUBLIC_KEY(MBEDTLS_TEST_PSA_ECC_ONE_FAMILY):MBEDTLS_TEST_PSA_ECC_ONE_CURVE_BITS': {
            'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_ECDH': [
                'PSA_KEY_USAGE_DERIVE_PUBLIC',
            ],
        },
        'PSA_KEY_TYPE_RSA_KEY_PAIR:1024': {
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_SIGN_HASH',
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_CRYPT': [
                'PSA_KEY_USAGE_DECRYPT',
                'PSA_KEY_USAGE_ENCRYPT',
            ],
            'PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_DECRYPT',
                'PSA_KEY_USAGE_ENCRYPT',
            ],
        },
        'PSA_KEY_TYPE_RSA_PUBLIC_KEY:1024': {
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_ANY_HASH)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PSS(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_VERIFY_HASH',
            ],
            'PSA_ALG_RSA_PKCS1V15_CRYPT': [
                'PSA_KEY_USAGE_ENCRYPT',
            ],
            'PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256)': [
                'PSA_KEY_USAGE_ENCRYPT',
            ],
        },
    }
}

# PSA_WANT dependencies for key types and alg
DEPENDS_ON = {
    'MBEDTLS_PK_ECDSA': ['PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY'],
    'MBEDTLS_PK_ECKEY': ['PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY'],
    'MBEDTLS_PK_ECKEY_DH': ['PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY'],
    'PSA_ALG_ECDSA': ['PSA_WANT_ALG_ECDSA'],
    'PSA_ALG_DETERMINISTIC_ECDSA': ['PSA_WANT_ALG_DETERMINISTIC_ECDSA'],
    'PSA_ALG_SHA_256': ['PSA_WANT_ALG_SHA_256'],
    'PSA_ALG_ECDH': ['PSA_WANT_ALG_ECDH'],
    'MBEDTLS_PK_RSA': ['PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY'],
    'PSA_ALG_RSA_PKCS1V15_SIGN': ['PSA_WANT_ALG_RSA_PKCS1V15_SIGN'],
    'PSA_ALG_RSA_PSS_ANY_SALT': ['PSA_WANT_ALG_RSA_PSS'],
    'PSA_ALG_RSA_PSS': ['PSA_WANT_ALG_RSA_PSS'],
    'PSA_ALG_RSA_PKCS1V15_CRYPT': ['PSA_WANT_ALG_RSA_PKCS1V15_CRYPT'],
    'PSA_ALG_RSA_OAEP': ['PSA_WANT_ALG_RSA_OAEP'],
    'MBEDTLS_TEST_PSA_ECC_ONE_FAMILY': ['MBEDTLS_TEST_PSA_ECC_AT_LEAST_ONE_CURVE'],
}

TEST_ALG = [
    # Signature
    'PSA_ALG_ECDSA(PSA_ALG_ANY_HASH)',
    'PSA_ALG_ECDSA(PSA_ALG_SHA_256)',
    'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH)',
    'PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256)',
    'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH)',
    'PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_SHA_256)',
    'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH)',
    'PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_SHA_256)',
    'PSA_ALG_RSA_PSS(PSA_ALG_ANY_HASH)',
    'PSA_ALG_RSA_PSS(PSA_ALG_SHA_256)',
    # Key agreement
    'PSA_ALG_ECDH',
    # Encryption
    'PSA_ALG_RSA_PKCS1V15_CRYPT',
    'PSA_ALG_RSA_OAEP(PSA_ALG_ANY_HASH)',
    'PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256)',
]

TEST_USAGE = [
    'PSA_KEY_USAGE_ENCRYPT',
    'PSA_KEY_USAGE_DECRYPT',
    'PSA_KEY_USAGE_SIGN_HASH',
    'PSA_KEY_USAGE_VERIFY_HASH',
    'PSA_KEY_USAGE_DERIVE',
    'PSA_KEY_USAGE_DERIVE_PUBLIC',
]

def populate_depends_on(**kwargs):
    """
    Generate the depends_on list based on the parameters that are passed. Check
    DEPENDS_ON dictionary for the full list.
    """
    depends_on = []
    known_deps = [dep
                  for arg in kwargs.values()
                  for dep in DEPENDS_ON
                  if dep in arg]
    for dep in known_deps:
        depends_on.extend(DEPENDS_ON[dep])
    # Usage of enrollment algorithm is only supported when PSA_CRYPTO_C is set.
    # This is important in psasim tests because in that case PSA_CRYPTO_C is
    # not set.
    if kwargs['alg2'] != 'PSA_ALG_NONE':
        depends_on.extend(['MBEDTLS_PSA_CRYPTO_C'])
    return depends_on

def alg_include_hash(alg):
    """
    Return 1 if the algorithm that's passed in includes some hashing; 0 otherwise
    """
    return 1 if ('PSA_ALG_ANY_HASH' in alg) or ('PSA_ALG_SHA_' in alg) else 0

def split_base_alg_from_hash(alg):
    """
    Separate the base algorithm and the hashing one and return them as touple:
    - 1st item is the base alg;
    - 2nd item is the hash alg or '' if the hash alg is not found.
    """
    base_alg = re.sub(r'\(.*\)', '', alg)
    match = re.search(r'\((.*)\)', alg)
    hash_alg = match.group(1) if match is not None else ''
    return (base_alg, hash_alg)

#pylint: disable=too-many-return-statements
def are_algorithms_compatible(test_alg, key_alg):
    """
    Helper to verify if key_alg can do test_alg when both include some hashing
    function.
    """
    test_alg_base, test_alg_hash = split_base_alg_from_hash(test_alg)
    key_alg_base, key_alg_hash = split_base_alg_from_hash(key_alg)

    if test_alg_base != key_alg_base:
        return 0

    if key_alg_hash in ('PSA_ALG_ANY_HASH', test_alg_hash):
        return 1

    return 0

# pylint: disable=too-many-arguments
def get_expected_result(pk_type, key_type_bits, alg, alg2, usage, test_usage, test_alg):
    """
    Get the expected result for the test case given all the parameters that were
    used to generate it.
    """

    # Transparent contexts with public keys have limited validation capabilities in
    # mbedtls_pk_can_do_psa() so we need a special handling here:
    if ('PUBLIC' in key_type_bits) and (pk_type != 'MBEDTLS_PK_OPAQUE'):
        if (pk_type == 'MBEDTLS_PK_ECDSA') and \
            ('ECDSA' in test_alg and 'VERIFY_HASH' in test_usage):
            return 1
        if (pk_type == 'MBEDTLS_PK_ECKEY') and \
            (('ECDSA' in test_alg and 'VERIFY_HASH' in test_usage) or \
             ('ECDH' in test_alg and 'DERIVE_PUBLIC' in test_usage)):
            return 1
        if (pk_type == 'MBEDTLS_PK_ECKEY_DH') and \
           ('ECDH' in test_alg and 'DERIVE_PUBLIC' in test_usage):
            return 1
        if (pk_type == 'MBEDTLS_PK_RSA') and \
           (('RSA_PKCS1V15_SIGN' in test_alg and 'VERIFY_HASH' in test_usage) or \
            ('RSA_PSS' in test_alg and 'VERIFY_HASH' in test_usage)):
            return 1
        if (pk_type == 'MBEDTLS_PK_RSA') and \
           (('RSA_PKCS1V15_CRYPT' in test_alg and 'ENCRYPT' in test_usage) or \
            ('RSA_OAEP' in test_alg and 'ENCRYPT' in test_usage)):
            return 1
        return 0

    # PSA_KEY_USAGE_DERIVE_PUBLIC deserves a special treatment (see the definition
    # of this symbol for further details). Therefore we only rely on key type and
    # testing algorithm to be OK.
    if ('PSA_KEY_USAGE_DERIVE_PUBLIC' in test_usage) and \
       ('ECC' in key_type_bits) and \
       ('PSA_ALG_ECDH' in test_alg):
        return 1

    # Check for wrong usage
    if test_usage != usage:
        return 0

    # Now handle transparent contexts with private key and opaque ones.
    # Check if algorithms match perfectly (simplest case).
    if test_alg in (alg, alg2):
        return 1

    # Check wildcard hash options
    if alg_include_hash(test_alg) and alg_include_hash(alg) and \
        are_algorithms_compatible(test_alg, alg):
        return 1
    if alg_include_hash(test_alg) and alg_include_hash(alg2) and \
        are_algorithms_compatible(test_alg, alg2):
        return 1

    return 0

def generate_test_case(test_number, pk_type, key_type_bits, alg, alg2, usage, test_alg, test_usage):
    """
    Generate a single test case to be written to the output data file.
    """
    new_content = []
    new_content.append(f"PK can do PSA #{test_number}: {pk_type} + {alg}/{alg2} + {usage} " \
                       f"-> {test_alg} + {test_usage}\n")
    depends_on = populate_depends_on(pk_type=pk_type, key_type_bits=key_type_bits, alg=alg, alg2=alg2)
    new_content.append("depends_on:" + ":".join(depends_on) + "\n")
    expected_result = get_expected_result(pk_type, key_type_bits, alg,
                                          alg2, usage, test_usage, test_alg)
    new_content.append(f"pk_can_do_psa:{test_number}:{pk_type}:{key_type_bits}:{alg}:{alg2}:" \
                       f"{usage}:{test_alg}:{test_usage}:{expected_result}\n")
    new_content.append("\n")

    return new_content

def main():
    test_number = 1
    output_content = []

    # Generate test cases
    # pylint: disable=too-many-nested-blocks
    for pk_type in PK_CONTEXTS:
        for key_type_bits in PK_CONTEXTS[pk_type]:
            for alg in PK_CONTEXTS[pk_type][key_type_bits]:
                for usage in PK_CONTEXTS[pk_type][key_type_bits][alg]:
                    for test_alg in TEST_ALG:
                        for test_usage in TEST_USAGE:
                            # Transparent contexts and opaque ones where enrollment alg is not set.
                            output_content.extend(
                                generate_test_case(test_number, pk_type, key_type_bits,
                                                   alg, 'PSA_ALG_NONE', usage,
                                                   test_alg, test_usage)
                            )
                            test_number = test_number + 1
                            # For opaque contexts check also the case in which alg is PSA_ALG_NONE
                            # and enrollment alg (i.e. alg2) is set.
                            if pk_type == 'MBEDTLS_PK_OPAQUE':
                                output_content.extend(
                                    generate_test_case(test_number, pk_type, key_type_bits,
                                                       'PSA_ALG_NONE', alg, usage,
                                                       test_alg, test_usage)
                                )
                                test_number = test_number + 1

    output_file = os.path.join(os.path.dirname(__file__),
                               os.path.pardir, os.path.pardir,
                               'tests', 'suites', 'test_suite_pk.pk_can_do_psa.data')

    with open(output_file, "wt") as fp:
        fp.write(HEADER_MESSAGE)
        fp.write("".join(output_content))

if __name__ == "__main__":
    main()
