Tools: add Mac0Message support to iatverifier

Add support for the Mac0Message COSE format to iatverifier, as the
alternative to the Sign1Message currently used.

Change-Id: I6baa87209fd17afe52ff1c6f936693e3b9dc9b9f
Signed-off-by: Sergei Trofimov <sergei.trofimov@arm.com>
diff --git a/tools/iat-verifier/README.rst b/tools/iat-verifier/README.rst
index dac808d..82def37 100644
--- a/tools/iat-verifier/README.rst
+++ b/tools/iat-verifier/README.rst
@@ -91,6 +91,22 @@
    }
 
 
+
+***********
+Mac0Message
+***********
+
+By default, the expectation is that the message will be wrapped using
+Sign1Message  COSE structure, however, the alternative Mac0Message structure
+that uses HMAC with SHA256 algorithm rather than a signature is supported via
+the ``-m mac`` flag:
+
+::
+
+    $ check_iat -m mac -k sample/hmac.key sample/iat-hmac.cbor
+    Signature OK
+    Token format OK
+
 *******
 Testing
 *******
@@ -126,4 +142,4 @@
 
 --------------
 
-*Copyright (c) 2019, Arm Limited. All rights reserved.*
+*Copyright (c) 2019-2020, Arm Limited. All rights reserved.*
diff --git a/tools/iat-verifier/iatverifier/util.py b/tools/iat-verifier/iatverifier/util.py
index 14a00cf..ffda7ac 100644
--- a/tools/iat-verifier/iatverifier/util.py
+++ b/tools/iat-verifier/iatverifier/util.py
@@ -1,5 +1,5 @@
 # -----------------------------------------------------------------------------
-# Copyright (c) 2019, Arm Limited. All rights reserved.
+# Copyright (c) 2019-2020, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -12,6 +12,7 @@
 import yaml
 from ecdsa import SigningKey, VerifyingKey
 from pycose.sign1message import Sign1Message
+from pycose.mac0message import Mac0Message
 
 from iatverifier import const
 
@@ -28,22 +29,39 @@
     return signed_msg.encode()
 
 
-def convert_map_to_token_files(mapfile, keyfile, outfile, raw=False):
+def hmac_eat(token, key=None):
+    hmac_msg = Mac0Message(payload=token, key=key)
+    hmac_msg.compute_auth_tag('HS256')
+    return hmac_msg.encode()
+
+
+def convert_map_to_token_files(mapfile, keyfile, outfile, method='sign'):
     token_map = read_token_map(mapfile)
 
-    with open(keyfile) as fh:
-        signing_key = SigningKey.from_pem(fh.read())
+    if method == 'sign':
+        with open(keyfile) as fh:
+            signing_key = SigningKey.from_pem(fh.read())
+    else:
+        with open(keyfile, 'rb') as fh:
+            signing_key = fh.read()
 
     with open(outfile, 'wb') as wfh:
         convert_map_to_token(token_map, signing_key, wfh, raw)
 
 
-def convert_map_to_token(token_map, signing_key, wfh, raw=False):
+def convert_map_to_token(token_map, signing_key, wfh, method='sign'):
     token = cbor.dumps(token_map)
-    if raw:
+
+    if method == 'raw':
         signed_token = token
-    else:
+    elif method == 'sign':
         signed_token = sign_eat(token, signing_key)
+    elif method == 'mac':
+        signed_token = hmac_eat(token, signing_key)
+    else:
+        err_msg = 'Unexpected method "{}"; must be one of: raw, sign, mac'
+        raise ValueError(err_msg.format(method))
+
     wfh.write(signed_token)
 
 
@@ -63,18 +81,28 @@
     return _parse_raw_token(raw)
 
 
-def extract_iat_from_cose(keyfile, tokenfile, keep_going=False):
-    key = read_keyfile(keyfile)
+def extract_iat_from_cose(keyfile, tokenfile, keep_going=False, method='sign'):
+    key = read_keyfile(keyfile, method)
 
     try:
         with open(tokenfile, 'rb') as wfh:
-            return get_cose_payload(wfh.read(), key)
+            return get_cose_payload(wfh.read(), key, method)
     except Exception as e:
         msg = 'Bad COSE file "{}": {}'
         raise ValueError(msg.format(tokenfile, e))
 
 
-def get_cose_payload(cose, key=None):
+def get_cose_payload(cose, key=None, method='sign'):
+    if method == 'sign':
+        return get_cose_sign1_payload(cose, key)
+    elif method == 'mac':
+        return get_cose_mac0_pyload(cose, key)
+    else:
+        err_msg = 'Unexpected method "{}"; must be one of: sign, mac'
+        raise ValueError(err_msg.format(method))
+
+
+def get_cose_sign1_payload(cose, key=None):
     msg = Sign1Message.decode(cose)
     if key:
         msg.key = key
@@ -86,6 +114,16 @@
     return msg.payload
 
 
+def get_cose_mac0_pyload(cose, key=None):
+    msg = Mac0Message.decode(cose)
+    if key:
+        msg.key = key
+        try:
+            msg.verify_auth_tag(alg='HS256')
+        except Exception as e:
+            raise ValueError('Bad signature ({})'.format(e))
+    return msg.payload
+
 def recursive_bytes_to_strings(d, in_place=False):
     if in_place:
         result = d
@@ -105,26 +143,40 @@
     return result
 
 
-def read_keyfile(keyfile):
+def read_keyfile(keyfile, method='sign'):
     if keyfile:
-        try:
-            key = SigningKey.from_pem(open(keyfile, 'rb').read())
-        except Exception as e:
-            signing_key_error = str(e)
-
-            try:
-                key = VerifyingKey.from_pem(open(keyfile, 'rb').read())
-            except Exception as e:
-                verifying_key_error = str(e)
-
-                msg = 'Bad key file "{}":\n\tpubkey error: {}\n\tprikey error: {}'
-                raise ValueError(msg.format(keyfile, verifying_key_error, signing_key_error))
+        if method == 'sign':
+            return read_sign1_key(keyfile)
+        elif method == 'mac':
+            return read_hmac_key(keyfile)
+        else:
+            err_msg = 'Unexpected method "{}"; must be one of: sign, mac'
+            raise ValueError(err_msg.format(method))
     else:  # no keyfile
         key = None
 
     return key
 
 
+def read_sign1_key(keyfile):
+    try:
+        key = SigningKey.from_pem(open(keyfile, 'rb').read())
+    except Exception as e:
+        signing_key_error = str(e)
+
+        try:
+            key = VerifyingKey.from_pem(open(keyfile, 'rb').read())
+        except Exception as e:
+            verifying_key_error = str(e)
+
+            msg = 'Bad key file "{}":\n\tpubkey error: {}\n\tprikey error: {}'
+            raise ValueError(msg.format(keyfile, verifying_key_error, signing_key_error))
+
+
+def read_hmac_key(keyfile):
+    return open(keyfile, 'rb').read()
+
+
 def _parse_raw_token(raw):
     result = {}
     for raw_key, raw_value in raw.items():
diff --git a/tools/iat-verifier/iatverifier/verify.py b/tools/iat-verifier/iatverifier/verify.py
index 13365c9..17d8b84 100644
--- a/tools/iat-verifier/iatverifier/verify.py
+++ b/tools/iat-verifier/iatverifier/verify.py
@@ -1,5 +1,5 @@
 # -----------------------------------------------------------------------------
-# Copyright (c) 2019, Arm Limited. All rights reserved.
+# Copyright (c) 2019-2020, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -289,13 +289,18 @@
                         help='''
                         Report failure if unknown claim is encountered.
                         ''')
+    parser.add_argument('-m', '--method', choices=['sign', 'mac'], default='sign',
+                        help='''
+                        Specify how this token is wrapped -- whether Sign1Message or
+                        Mac0Message COSE structure is used.
+                        ''')
     args = parser.parse_args()
 
     logging.basicConfig(level=logging.INFO)
 
     try:
         raw_iat = extract_iat_from_cose(args.keyfile, args.tokenfile,
-                                        args.keep_going)
+                                        args.keep_going, args.method)
         if args.keyfile and not seen_errors:
             print('Signature OK')
     except ValueError as e:
diff --git a/tools/iat-verifier/sample/hmac.key b/tools/iat-verifier/sample/hmac.key
new file mode 100644
index 0000000..e9569f1
--- /dev/null
+++ b/tools/iat-verifier/sample/hmac.key
Binary files differ
diff --git a/tools/iat-verifier/sample/iat-hmac.cbor b/tools/iat-verifier/sample/iat-hmac.cbor
new file mode 100644
index 0000000..1ea3018
--- /dev/null
+++ b/tools/iat-verifier/sample/iat-hmac.cbor
Binary files differ
diff --git a/tools/iat-verifier/scripts/compile_token b/tools/iat-verifier/scripts/compile_token
index 76e6e33..8d00f52 100755
--- a/tools/iat-verifier/scripts/compile_token
+++ b/tools/iat-verifier/scripts/compile_token
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #-------------------------------------------------------------------------------
-# Copyright (c) 2019, Arm Limited. All rights reserved.
+# Copyright (c) 2019-2020, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -21,28 +21,43 @@
     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('-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.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
+
+    if args.hmac:
+        method = 'hmac'
+        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 = 'raw'
+    else:
+        method = 'sign'
+        if args.keyfile:
+            with open(args.keyfile) as fh:
+                signing_key = SigningKey.from_pem(fh.read())
 
     token_map = read_token_map(args.source)
-    if args.keyfile:
-        with open(args.keyfile) as fh:
-            signing_key = SigningKey.from_pem(fh.read())
-    else:
-        signing_key = None
 
     if args.outfile:
         with open(args.outfile, 'wb') as wfh:
-            convert_map_to_token(token_map, signing_key, wfh, args.raw)
+            convert_map_to_token(token_map, signing_key, wfh, method)
     else:
         with os.fdopen(sys.stdout.fileno(), 'wb') as wfh:
-            convert_map_to_token(token_map, signing_key, wfh, args.raw)
+            convert_map_to_token(token_map, signing_key, wfh, method)
 
 
diff --git a/tools/iat-verifier/tests/data/hmac.key b/tools/iat-verifier/tests/data/hmac.key
new file mode 100644
index 0000000..e9569f1
--- /dev/null
+++ b/tools/iat-verifier/tests/data/hmac.key
Binary files differ