From 3703f2c20008936b1de85f1db7060884edf3a67e Mon Sep 17 00:00:00 2001
From: Dain Nilsson <dain@yubico.com>
Date: Wed, 23 Oct 2024 15:25:37 +0200
Subject: [PATCH] Add ARKG derived key class

---
 examples/sign_arkg.py     |  6 +++-
 fido2/arkg.py             | 71 ++++++++++++++++++++++++++++-----------
 fido2/cose.py             |  6 ++--
 fido2/ctap2/extensions.py |  2 +-
 4 files changed, 60 insertions(+), 25 deletions(-)

diff --git a/examples/sign_arkg.py b/examples/sign_arkg.py
index ab732a99..7e998a08 100644
--- a/examples/sign_arkg.py
+++ b/examples/sign_arkg.py
@@ -138,7 +138,11 @@ def request_uv(self, permissions, rd_id):
 # Arbitrary bytestring used for info
 info = b"my-info-here"
 # Derived public key to verify with, and kh to send to Authenticator
-pk2, kh = pk.derive_public_key(info)
+pk2 = pk.derive_public_key(info)
+print("Derived public key", pk2)
+ref = pk2.get_ref()
+print("COSE Key ref for derived key", ref)
+kh = cbor.encode(ref)
 
 # Prepare a message to sign
 message = b"New message"
diff --git a/fido2/arkg.py b/fido2/arkg.py
index e5ad238f..0eb9480b 100644
--- a/fido2/arkg.py
+++ b/fido2/arkg.py
@@ -1,6 +1,32 @@
-from .cose import CoseKey
+# Copyright (c) 2024 Yubico AB
+# All rights reserved.
+#
+#   Redistribution and use in source and binary forms, with or
+#   without modification, are permitted provided that the following
+#   conditions are met:
+#
+#    1. Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#    2. Redistributions in binary form must reproduce the above
+#       copyright notice, this list of conditions and the following
+#       disclaimer in the documentation and/or other materials provided
+#       with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+from .cose import CoseKey, ES256
 from .utils import int2bytes, bytes2int
-from . import cbor
 
 from cryptography.hazmat.primitives.hashes import SHA256, Hash, HashAlgorithm
 from cryptography.hazmat.primitives.hmac import HMAC
@@ -337,9 +363,25 @@ def _cose2point(cose):
     return SEC1Encoder.decode_public_key(b"\x04" + cose[-2] + cose[-3], P256)
 
 
+class ARKG_P256_DERIVED(ES256):
+    def __init__(self, *args, parent_kid: bytes, kh: bytes, info: bytes, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._parent_kid = parent_kid
+        self._kh = kh
+        self._info = info
+
+    def get_ref(self):
+        return {
+            1: -65538,  # kty: Ref-ARKG-derived
+            2: self._parent_kid,
+            3: ARKG_P256ADD_ECDH.ALGORITHM,
+            -1: self._kh,
+            -2: self._info,
+        }
+
+
 class ARKG_P256ADD_ECDH(CoseKey):
     ALGORITHM = -65539
-    _HASH_ALG = SHA256()
     _ARKG = ARKG(
         bl=BL(crv=P256, Hash=SHA256, DST_ext=b"ARKG-P256ADD-ECDH"),
         kem=KEM(crv=P256, Hash=SHA256, DST_ext=b"ARKG-P256ADD-ECDH"),
@@ -353,32 +395,21 @@ def blinding_key(self) -> CoseKey:
     def kem_key(self) -> CoseKey:
         return CoseKey.parse(self[-2])
 
-    def derive_public_key(self, info: bytes) -> Tuple[CoseKey, bytes]:
+    def derive_public_key(self, info: bytes) -> CoseKey:
         point, kh = self._ARKG.derive_public_key(
             _cose2point(self.kem_key),
             _cose2point(self.blinding_key),
             info,
         )
-        pk = CoseKey.parse(
+        return ARKG_P256_DERIVED(
             {
                 1: 2,
                 3: -7,
                 -1: 1,
                 -2: int2bytes(point.x, 32),
                 -3: int2bytes(point.y, 32),
-            }
-        )
-
-        # Return the complete COSE key ref used as key handle to getAssertion
-        ref = dict(self.get_key_handle())
-        ref.update(
-            {
-                1: -65538,  # kty: Ref-ARKG-derived
-                -1: kh,
-                -2: info,
-            }
+            },
+            parent_kid=self[2],
+            kh=kh,
+            info=info,
         )
-
-        ref_kh = cbor.encode(ref)
-
-        return pk, ref_kh
diff --git a/fido2/cose.py b/fido2/cose.py
index 90965203..75eff394 100644
--- a/fido2/cose.py
+++ b/fido2/cose.py
@@ -43,7 +43,7 @@ class CoseKey(dict):
 
     ALGORITHM: int = None  # type: ignore
 
-    def get_key_handle(self) -> Mapping[int, Any]:
+    def get_ref(self) -> Mapping[int, Any]:
         """Returns a COSE Key Reference for the key."""
         return {k: self[k] for k in (1, 2, 3) if k in self}
 
@@ -125,8 +125,8 @@ class ES256(CoseKey):
     ALGORITHM = -7
     _HASH_ALG = hashes.SHA256()
 
-    def get_key_handle(self):
-        kh = dict(super().get_key_handle())
+    def get_ref(self):
+        kh = dict(super().get_ref())
         kh[1] = -2  # Ref-EC2
         return kh
 
diff --git a/fido2/ctap2/extensions.py b/fido2/ctap2/extensions.py
index 480e4832..5b68787a 100644
--- a/fido2/ctap2/extensions.py
+++ b/fido2/ctap2/extensions.py
@@ -481,7 +481,7 @@ def process_create_output(self, attestation_response, *args):
         output = {
             "generatedKey": {
                 "publicKey": cbor.encode(pk),
-                "keyHandle": cbor.encode(pk.get_key_handle()),
+                "keyHandle": cbor.encode(pk.get_ref()),
             }
         }