feat: added support for EC2Key

IdentityPython · Jan 31, 2024 · 6935298 · 6935298
1 parent 0be45af
commit 6935298
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/pymdoccbor/mdoc/issuer.py b/pymdoccbor/mdoc/issuer.py
@@ -2,7 +2,7 @@
 import cbor2
 import logging
 
-from pycose.keys import CoseKey
+from pycose.keys import CoseKey, EC2Key
 from typing import Union
 
 from pymdoccbor.mso.issuer import MsoIssuer
@@ -16,7 +16,7 @@ class MdocCborIssuer:
     MdocCborIssuer helper class to create a new mdoc
     """
 
-    def __init__(self, private_key: Union[dict, CoseKey]):
+    def __init__(self, private_key: Union[dict, EC2Key, CoseKey]):
         """
         Create a new MdocCborIssuer instance
 
@@ -28,11 +28,17 @@ def __init__(self, private_key: Union[dict, CoseKey]):
         self.version: str = '1.0'
         self.status: int = 0
 
-        if not private_key:
+        if isinstance(private_key, dict):
+            self.private_key = CoseKey.from_dict(private_key)
+        elif isinstance(private_key, EC2Key):
+            ec2_encoded = private_key.encode()
+            ec2_decoded = CoseKey.decode(ec2_encoded)
+            self.private_key = ec2_decoded
+        elif isinstance(private_key, CoseKey):
+            self.private_key = private_key
+        else:
             raise MissingPrivateKey("You must provide a private key")
 
-        if private_key and isinstance(private_key, dict):
-            self.private_key = CoseKey.from_dict(private_key)
 
         self.signed :dict = {}
 

diff --git a/pymdoccbor/mso/issuer.py b/pymdoccbor/mso/issuer.py
@@ -26,7 +26,7 @@ class MsoIssuer(MsoX509Fabric):
     def __init__(
         self,
         data: dict,
-        private_key: Union[dict, CoseKey],
+        private_key: Union[dict, EC2Key, CoseKey],
         digest_alg: str = settings.PYMDOC_HASHALG
     ):
         """
@@ -48,6 +48,10 @@ def __init__(
                 self.private_key.kid = str(uuid.uuid4())
         elif private_key and isinstance(private_key, CoseKey):
             self.private_key = private_key
+        elif private_key and isinstance(private_key, EC2Key):
+            ec2_encoded = private_key.encode()
+            ec2_decoded = CoseKey.decode(ec2_encoded)
+            self.private_key = ec2_decoded
         else:
             raise MsoPrivateKeyRequired(
                 "MSO Writer requires a valid private key"