Use enums instead of literals

Using enums makes it easier to use the functions in a typesafe way and
to validate user input.

Fixes: #69

Author: robin-nitrokey

nethsm/__init__.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from io import BufferedReader
-from typing import TYPE_CHECKING, Any, Iterator, Literal, Mapping, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Union, cast
 from urllib.parse import urlencode
 
 import urllib3
@@ -76,6 +76,13 @@ class UnattendedBootStatus(enum.Enum):
     ON = "on"
     OFF = "off"
 
+    @staticmethod
+    def from_string(s: str) -> "UnattendedBootStatus":
+        for status in UnattendedBootStatus:
+            if status.value == s:
+                return status
+        raise ValueError(f"Unsupported unattended boot status {s}")
+
 
 class KeyType(enum.Enum):
     RSA = "RSA"
@@ -94,11 +101,6 @@ def from_string(s: str) -> "KeyType":
         raise ValueError(f"Unsupported key type {s}")
 
 
-KeyTypeLitteral = Literal[
-    "RSA", "Curve25519", "EC_P224", "EC_P256", "EC_P384", "EC_P521", "Generic"
-]
-
-
 class KeyMechanism(enum.Enum):
     RSA_DECRYPTION_RAW = "RSA_Decryption_RAW"
     RSA_DECRYPTION_PKCS1 = "RSA_Decryption_PKCS1"
@@ -120,33 +122,24 @@ class KeyMechanism(enum.Enum):
     AES_ENCRYPTION_CBC = "AES_Encryption_CBC"
     AES_DECRYPTION_CBC = "AES_Decryption_CBC"
 
-
-KeyMechanismLiteral = Literal[
-    "RSA_Decryption_RAW",
-    "RSA_Decryption_PKCS1",
-    "RSA_Decryption_OAEP_MD5",
-    "RSA_Decryption_OAEP_SHA1",
-    "RSA_Decryption_OAEP_SHA224",
-    "RSA_Decryption_OAEP_SHA256",
-    "RSA_Decryption_OAEP_SHA384",
-    "RSA_Decryption_OAEP_SHA512",
-    "RSA_Signature_PKCS1",
-    "RSA_Signature_PSS_MD5",
-    "RSA_Signature_PSS_SHA1",
-    "RSA_Signature_PSS_SHA224",
-    "RSA_Signature_PSS_SHA256",
-    "RSA_Signature_PSS_SHA384",
-    "RSA_Signature_PSS_SHA512",
-    "EdDSA_Signature",
-    "ECDSA_Signature",
-    "AES_Encryption_CBC",
-    "AES_Decryption_CBC",
-]
+    @staticmethod
+    def from_string(s: str) -> "KeyMechanism":
+        for key_mechanism in KeyMechanism:
+            if key_mechanism.value == s:
+                return key_mechanism
+        raise ValueError(f"Unsupported key mechanism {s}")
 
 
 class EncryptMode(enum.Enum):
     AES_CBC = "AES_CBC"
 
+    @staticmethod
+    def from_string(s: str) -> "EncryptMode":
+        for mode in EncryptMode:
+            if mode.value == s:
+                return mode
+        raise ValueError(f"Unsupported encrypt mode {s}")
+
 
 class DecryptMode(enum.Enum):
     RAW = "RAW"
@@ -159,6 +152,13 @@ class DecryptMode(enum.Enum):
     OAEP_SHA512 = "OAEP_SHA512"
     AES_CBC = "AES_CBC"
 
+    @staticmethod
+    def from_string(s: str) -> "DecryptMode":
+        for mode in DecryptMode:
+            if mode.value == s:
+                return mode
+        raise ValueError(f"Unsupported decrypt mode {s}")
+
 
 class SignMode(enum.Enum):
     PKCS1 = "PKCS1"
@@ -171,6 +171,13 @@ class SignMode(enum.Enum):
     EDDSA = "EdDSA"
     ECDSA = "ECDSA"
 
+    @staticmethod
+    def from_string(s: str) -> "SignMode":
+        for mode in SignMode:
+            if mode.value == s:
+                return mode
+        raise ValueError(f"Unsupported sign mode {s}")
+
 
 class TlsKeyType(enum.Enum):
     RSA = "RSA"
@@ -180,6 +187,13 @@ class TlsKeyType(enum.Enum):
     EC_P384 = "EC_P384"
     EC_P521 = "EC_P521"
 
+    @staticmethod
+    def from_string(s: str) -> "TlsKeyType":
+        for key_type in TlsKeyType:
+            if key_type.value == s:
+                return key_type
+        raise ValueError(f"Unsupported TLS key type {s}")
+
 
 @dataclass
 class SystemInfo:
@@ -199,7 +213,7 @@ class User:
 @dataclass
 class Key:
     key_id: str
-    mechanisms: list[str]
+    mechanisms: list[KeyMechanism]
     type: KeyType
     operations: int
     tags: Optional[list[str]]
@@ -521,7 +535,7 @@ def get_user(self, user_id: str) -> User:
     def add_user(
         self,
         real_name: str,
-        role: Literal["Administrator", "Operator", "Metrics", "Backup"],
+        role: Role,
         passphrase: str,
         user_id: Optional[str] = None,
     ) -> str:
@@ -530,7 +544,7 @@ def add_user(
 
         body = UserPostDataDict(
             realName=real_name,
-            role=role,
+            role=role.value,
             passphrase=passphrase,
         )
         try:
@@ -769,7 +783,9 @@ def get_key(self, key_id: str) -> Key:
             )
         return Key(
             key_id=key_id,
-            mechanisms=[mechanism for mechanism in key.mechanisms],
+            mechanisms=[
+                KeyMechanism.from_string(mechanism) for mechanism in key.mechanisms
+            ],
             type=KeyType.from_string(key.type),
             operations=key.operations,
             tags=[str(tag) for tag in cast(list[str], key.restrictions["tags"])]
@@ -805,14 +821,15 @@ def get_key_public_key(self, key_id: str) -> str:
     def add_key(
         self,
         key_id: str,
-        type: KeyTypeLitteral,
-        mechanisms: list[KeyMechanismLiteral],
+        type: KeyType,
+        mechanisms: list[KeyMechanism],
         tags: list[str],
         prime_p: Optional[str],
         prime_q: Optional[str],
         public_exponent: Optional[str],
         data: Optional[str],
     ) -> str:
+        from .client.components.schema.key_mechanisms import KeyMechanismsTupleInput
         from .client.components.schema.key_private_data import KeyPrivateDataDict
         from .client.components.schema.key_restrictions import KeyRestrictionsDict
         from .client.components.schema.private_key import PrivateKeyDict
@@ -821,7 +838,7 @@ def add_key(
         # To do: split into different methods for RSA and other key types, or
         # at least change typing accordingly
 
-        if type == "RSA":
+        if type == KeyType.RSA:
             assert prime_p
             assert prime_q
             assert public_exponent
@@ -834,19 +851,23 @@ def add_key(
             assert data
             key_data = KeyPrivateDataDict(data=data)
 
+        mechanism_tuple: KeyMechanismsTupleInput = [
+            mechanism.value for mechanism in mechanisms
+        ]
+
         if tags:
             body = PrivateKeyDict(
-                type=type,
-                mechanisms=mechanisms,
+                type=type.value,
+                mechanisms=mechanism_tuple,
                 key=key_data,
                 restrictions=KeyRestrictionsDict(
                     tags=TagListTuple([tag for tag in tags])
                 ),
             )
         else:
             body = PrivateKeyDict(
-                type=type,
-                mechanisms=mechanisms,
+                type=type.value,
+                mechanisms=mechanism_tuple,
                 key=key_data,
             )
 
@@ -899,27 +920,31 @@ def delete_key(self, key_id: str) -> None:
 
     def generate_key(
         self,
-        type: KeyTypeLitteral,
-        mechanisms: tuple[KeyMechanismLiteral],
+        type: KeyType,
+        mechanisms: list[KeyMechanism],
         length: int,
         key_id: Optional[str] = None,
     ) -> str:
         from .client.components.schema.key_generate_request_data import (
             KeyGenerateRequestDataDict,
         )
-        from .client.components.schema.key_mechanisms import KeyMechanismsTuple
+        from .client.components.schema.key_mechanisms import KeyMechanismsTupleInput
+
+        mechanism_tuple: KeyMechanismsTupleInput = [
+            mechanism.value for mechanism in mechanisms
+        ]
 
         if key_id:
             body = KeyGenerateRequestDataDict(
-                type=type,
-                mechanisms=KeyMechanismsTuple(mechanisms),
+                type=type.value,
+                mechanisms=mechanism_tuple,
                 length=length,
                 id=key_id,
             )
         else:
             body = KeyGenerateRequestDataDict(
-                type=type,
-                mechanisms=KeyMechanismsTuple(mechanisms),
+                type=type.value,
+                mechanisms=mechanism_tuple,
                 length=length,
             )
         try:
@@ -1101,7 +1126,7 @@ def csr(
 
     def generate_tls_key(
         self,
-        type: Literal["RSA", "Curve25519", "EC_P224", "EC_P256", "EC_P384", "EC_P521"],
+        type: TlsKeyType,
         length: Optional[int] = None,
     ) -> None:
         from .client.components.schema.tls_key_generate_request_data import (
@@ -1110,7 +1135,7 @@ def generate_tls_key(
         from .client.schemas import Unset
 
         body = TlsKeyGenerateRequestDataDict(
-            type=type,
+            type=type.value,
             length=length if length is not None else Unset(),
         )
 
@@ -1207,11 +1232,13 @@ def set_logging_config(
         self,
         ip_address: str,
         port: int,
-        log_level: Literal["debug", "info", "warning", "error"],
+        log_level: LogLevel,
     ) -> None:
         from .client.components.schema.logging_config import LoggingConfigDict
 
-        body = LoggingConfigDict(ipAddress=ip_address, port=port, logLevel=log_level)
+        body = LoggingConfigDict(
+            ipAddress=ip_address, port=port, logLevel=log_level.value
+        )
         try:
             self.get_api().config_logging_put(body=body)
         except Exception as e:
@@ -1256,12 +1283,12 @@ def set_time(self, time: Union[str, datetime]) -> None:
                 },
             )
 
-    def set_unattended_boot(self, status: Literal["on", "off"]) -> None:
+    def set_unattended_boot(self, status: UnattendedBootStatus) -> None:
         from .client.components.schema.unattended_boot_config import (
             UnattendedBootConfigDict,
         )
 
-        body = UnattendedBootConfigDict(status=status)
+        body = UnattendedBootConfigDict(status=status.value)
         try:
             self.get_api().config_unattended_boot_put(body=body)
         except Exception as e:
@@ -1402,7 +1429,7 @@ def factory_reset(self) -> None:
             )
 
     def encrypt(
-        self, key_id: str, data: str, mode: Literal["AES_CBC"], iv: str
+        self, key_id: str, data: str, mode: EncryptMode, iv: str
     ) -> tuple[str, str]:
         from .client.components.schema.encrypt_request_data import (
             EncryptRequestDataDict,
@@ -1412,7 +1439,7 @@ def encrypt(
         )
 
         path_params = PathParametersDict(KeyID=key_id)
-        body = EncryptRequestDataDict(message=data, mode=mode, iv=iv)
+        body = EncryptRequestDataDict(message=data, mode=mode.value, iv=iv)
         try:
             response = self.get_api().keys_key_id_encrypt_post(
                 path_params=path_params, body=body
@@ -1433,17 +1460,7 @@ def decrypt(
         self,
         key_id: str,
         data: str,
-        mode: Literal[
-            "RAW",
-            "PKCS1",
-            "OAEP_MD5",
-            "OAEP_SHA1",
-            "OAEP_SHA224",
-            "OAEP_SHA256",
-            "OAEP_SHA384",
-            "OAEP_SHA512",
-            "AES_CBC",
-        ],
+        mode: DecryptMode,
         iv: str,
     ) -> str:
         from .client.components.schema.decrypt_request_data import (
@@ -1453,10 +1470,10 @@ def decrypt(
             PathParametersDict,
         )
 
-        body = DecryptRequestDataDict(encrypted=data, mode=mode, iv=iv)
+        body = DecryptRequestDataDict(encrypted=data, mode=mode.value, iv=iv)
 
         if len(iv) == 0:
-            body = DecryptRequestDataDict(encrypted=data, mode=mode)
+            body = DecryptRequestDataDict(encrypted=data, mode=mode.value)
 
         path_params = PathParametersDict(KeyID=key_id)
         try:
@@ -1479,25 +1496,15 @@ def sign(
         self,
         key_id: str,
         data: str,
-        mode: Literal[
-            "PKCS1",
-            "PSS_MD5",
-            "PSS_SHA1",
-            "PSS_SHA224",
-            "PSS_SHA256",
-            "PSS_SHA384",
-            "PSS_SHA512",
-            "EdDSA",
-            "ECDSA",
-        ],
+        mode: SignMode,
     ) -> str:
         from .client.components.schema.sign_request_data import SignRequestDataDict
         from .client.paths.keys_key_id_sign.post.path_parameters import (
             PathParametersDict,
         )
 
         path_params = PathParametersDict(KeyID=key_id)
-        body = SignRequestDataDict(message=data, mode=mode)
+        body = SignRequestDataDict(message=data, mode=mode.value)
         try:
             response = self.get_api().keys_key_id_sign_post(
                 path_params=path_params, body=body