diff --git a/jrnl/encryption/__init__.py b/jrnl/encryption/__init__.py index 7fbb8e46..aed00297 100644 --- a/jrnl/encryption/__init__.py +++ b/jrnl/encryption/__init__.py @@ -1,22 +1,31 @@ # Copyright © 2012-2022 jrnl contributors # License: https://www.gnu.org/licenses/gpl-3.0.html +from enum import Enum from importlib import import_module +from .BaseEncryption import BaseEncryption + + +class EncryptionMethods(str, Enum): + NONE = "NoEncryption" + JRNLV1 = "Jrnlv1Encryption" + JRNLV2 = "Jrnlv2Encryption" + + ENCRYPTION_METHODS = { # config: classname (as string) - "default": "NoEncryption", - False: "NoEncryption", - True: "Jrnlv2Encryption", - "jrnlv1": "Jrnlv1Encryption", - "jrnlv2": "Jrnlv2Encryption", + True: EncryptionMethods.JRNLV2, # the default + False: EncryptionMethods.NONE, + "jrnlv1": EncryptionMethods.JRNLV1, + "jrnlv2": EncryptionMethods.JRNLV2, } -def determine_encryption_method(config): +def determine_encryption_method(config: str | bool) -> BaseEncryption: key = config if isinstance(config, str): key = config.lower() - my_class = ENCRYPTION_METHODS.get(key, "default") + my_class = ENCRYPTION_METHODS[key] return getattr(import_module(f"jrnl.encryption.{my_class}"), my_class)