Refactor get_key and Key

nethsm/__init__.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from io import BufferedReader, FileIO
-from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Union
 from urllib.parse import urlencode
 
 import urllib3
@@ -219,16 +219,28 @@ class User:
     role: Role
 
 
+@dataclass
+class RsaPublicKey:
+    modulus: str
+    public_exponent: str
+
+
+@dataclass
+class EcPublicKey:
+    data: str
+
+
+PublicKey = Union[RsaPublicKey, EcPublicKey, None]
+
+
 @dataclass
 class Key:
     key_id: str
     mechanisms: list[KeyMechanism]
     type: KeyType
     operations: int
-    tags: Optional[list[str]]
-    modulus: Optional[str]
-    public_exponent: Optional[str]
-    data: Optional[str]
+    tags: list[str]
+    public_key: PublicKey
 
 
 @dataclass
@@ -787,6 +799,7 @@ def list_keys(self, filter: Optional[str] = None) -> list[str]:
 
     def get_key(self, key_id: str) -> Key:
         from .client.paths.keys_key_id.get.path_parameters import PathParametersDict
+        from .client.schemas import Unset
 
         path_params = PathParametersDict(KeyID=key_id)
         try:
@@ -801,19 +814,45 @@ def get_key(self, key_id: str) -> Key:
                     404: f"Key {key_id} not found",
                 },
             )
+
+        mechanisms = [
+            KeyMechanism.from_string(mechanism) for mechanism in key.mechanisms
+        ]
+        tags = []
+        if not isinstance(key.restrictions, Unset):
+            if not isinstance(key.restrictions.tags, Unset):
+                tags = list(key.restrictions.tags)
+        key_type = KeyType.from_string(key.type)
+
+        public_key: PublicKey
+        if key_type == KeyType.RSA:
+            assert not isinstance(key.key, Unset)
+            assert isinstance(key.key.data, Unset)
+            assert not isinstance(key.key.modulus, Unset)
+            assert not isinstance(key.key.publicExponent, Unset)
+            public_key = RsaPublicKey(
+                modulus=key.key.modulus, public_exponent=key.key.publicExponent
+            )
+        elif key_type == KeyType.GENERIC:
+            if not isinstance(key.key, Unset):
+                assert isinstance(key.key.data, Unset)
+                assert isinstance(key.key.modulus, Unset)
+                assert isinstance(key.key.publicExponent, Unset)
+            public_key = None
+        else:
+            assert not isinstance(key.key, Unset)
+            assert not isinstance(key.key.data, Unset)
+            assert isinstance(key.key.modulus, Unset)
+            assert isinstance(key.key.publicExponent, Unset)
+            public_key = EcPublicKey(data=key.key.data)
+
         return Key(
             key_id=key_id,
-            mechanisms=[
-                KeyMechanism.from_string(mechanism) for mechanism in key.mechanisms
-            ],
-            type=KeyType.from_string(key.type),
+            mechanisms=mechanisms,
+            type=key_type,
             operations=key.operations,
-            tags=[str(tag) for tag in cast(list[str], key.restrictions["tags"])]
-            if "tags" in key.restrictions.keys()
-            else None,
-            modulus=getattr(key.key, "modulus", None),
-            public_exponent=getattr(key.key, "public_exponent", None),
-            data=getattr(key.key, "data", None),
+            tags=tags,
+            public_key=public_key,
         )
 
     # Get the public key file in PEM format

tests/test_nethsm_keys.py

@@ -15,7 +15,15 @@
 )
 
 import nethsm as nethsm_module
-from nethsm import DecryptMode, EncryptMode, KeyMechanism, KeyType, NetHSM, SignMode
+from nethsm import (
+    DecryptMode,
+    EncryptMode,
+    KeyMechanism,
+    KeyType,
+    NetHSM,
+    RsaPublicKey,
+    SignMode,
+)
 
 """########## Preparation for the Tests ##########
 
@@ -141,12 +149,9 @@ def test_generate_get_key_by_id(nethsm: nethsm_module.NetHSM) -> None:
     for mechanism in key.mechanisms:
         assert mechanism in C.MECHANISM
     assert key.operations >= 0
-    if key.tags:
-        assert key.tags
-    if key.modulus:
-        assert key.modulus
-    if key.public_exponent:
-        assert key.public_exponent
+    assert isinstance(key.public_key, RsaPublicKey)
+    assert key.public_key.modulus
+    assert key.public_key.public_exponent
 
 
 def test_add_key_tag_get_key(nethsm: NetHSM) -> None:
@@ -159,7 +164,6 @@ def test_add_key_tag_get_key(nethsm: NetHSM) -> None:
 
     key = nethsm.get_key(C.KEY_ID_GENERATED)
     tags = key.tags
-    assert tags
     assert C.TAG1 in tags
     assert C.TAG2 in tags
     assert C.TAG3 in tags
@@ -177,7 +181,6 @@ def test_delete_key_tag_get_key(nethsm: NetHSM) -> None:
     nethsm.delete_key_tag(key_id=C.KEY_ID_GENERATED, tag=C.TAG2)
     key = nethsm.get_key(C.KEY_ID_GENERATED)
 
-    assert key.tags
     assert C.TAG1 not in key.tags
     assert C.TAG2 not in key.tags
     assert C.TAG3 in key.tags
@@ -198,12 +201,9 @@ def test_list_get_keys(nethsm: nethsm_module.NetHSM) -> None:
         for mechanism in key.mechanisms:
             assert mechanism in C.MECHANISM
         assert key.operations >= 0
-        if key.tags:
-            assert key.tags
-        if key.modulus:
-            assert key.modulus
-        if key.public_exponent:
-            assert key.public_exponent
+        assert isinstance(key.public_key, RsaPublicKey)
+        assert key.public_key.modulus
+        assert key.public_key.public_exponent
 
 
 def test_delete_key(nethsm: NetHSM) -> None: