Add command line option for selecting token type

Change-Id: Id4b18d34d0897033490e785706cebe9070606b6c
Signed-off-by: Mate Toth-Pal <mate.toth-pal@arm.com>
diff --git a/iat-verifier/iatverifier/verify.py b/iat-verifier/iatverifier/verify.py
index eab1ea6..97c8078 100644
--- a/iat-verifier/iatverifier/verify.py
+++ b/iat-verifier/iatverifier/verify.py
@@ -18,6 +18,11 @@
 logger = logging.getLogger('iat-verify')
 
 def main():
+
+    token_verifiers = {
+        "PSA-IoT-Profile1-token": PSAIoTProfile1TokenVerifier,
+    }
+
     parser = argparse.ArgumentParser(
         description='''
         Validates a signed Initial Attestation Token (IAT), checking
@@ -49,12 +54,17 @@
                         Specify how this token is wrapped -- whether Sign1Message or
                         Mac0Message COSE structure is used.
                         ''')
+    parser.add_argument('-t', '--token-type',
+                        help='''The type of the Token.''',
+                        choices=token_verifiers.keys(),
+                        required=True)
+
     args = parser.parse_args()
 
     logging.basicConfig(level=logging.INFO)
 
     config = VerifierConfiguration(keep_going=args.keep_going, strict=args.strict)
-    verifier = PSAIoTProfile1TokenVerifier.get_verifier(config)
+    verifier = token_verifiers[args.token_type].get_verifier(config)
     if args.method == 'mac':
         verifier.method = AttestationTokenVerifier.SIGN_METHOD_MAC0
         verifier.cose_alg = AttestationTokenVerifier.COSE_ALG_HS256
diff --git a/iat-verifier/scripts/compile_token b/iat-verifier/scripts/compile_token
index 609e2c2..22fbe23 100755
--- a/iat-verifier/scripts/compile_token
+++ b/iat-verifier/scripts/compile_token
@@ -20,6 +20,10 @@
 if __name__ == '__main__':
     logging.basicConfig(level=logging.INFO)
 
+    token_verifiers = {
+        "PSA-IoT-Profile1-token": PSAIoTProfile1TokenVerifier,
+    }
+
     parser = argparse.ArgumentParser()
     parser.add_argument('source', help='Token source in YAML format')
     parser.add_argument('-o', '--outfile',
@@ -36,6 +40,10 @@
     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()
     signing_key = None
@@ -58,7 +66,7 @@
             with open(args.keyfile) as fh:
                 signing_key = SigningKey.from_pem(fh.read())
 
-    verifier = PSAIoTProfile1TokenVerifier.get_verifier()
+    verifier = token_verifiers[args.token_type].get_verifier()
     if verifier.method != method:
         verifier.method = method
     if cose_alg is not None and verifier.cose_alg != cose_alg:
diff --git a/iat-verifier/scripts/decompile_token b/iat-verifier/scripts/decompile_token
index cc74e12..d61247f 100755
--- a/iat-verifier/scripts/decompile_token
+++ b/iat-verifier/scripts/decompile_token
@@ -15,14 +15,23 @@
 
 
 if __name__ == '__main__':
+
+    token_verifiers = {
+        "PSA-IoT-Profile1-token": PSAIoTProfile1TokenVerifier,
+    }
+
     parser = argparse.ArgumentParser()
     parser.add_argument('source', help='A compiled COSE IAT token.')
     parser.add_argument('-o', '--outfile',
                         help='''Output file for the depompiled claims. If this is not
                         specified, the claims will be written to standard output.''')
+    parser.add_argument('-t', '--token-type',
+                        help='''The type of the Token.''',
+                        choices=token_verifiers.keys(),
+                        required=True)
     args = parser.parse_args()
 
-    verifier = PSAIoTProfile1TokenVerifier.get_verifier()
+    verifier = token_verifiers[args.token_type].get_verifier()
     with open(args.source, 'rb') as fh:
         token_map = convert_token_to_map(fh.read(), verifier)