From a9c78272ad280c57cf42346d8c3659cef15d720f Mon Sep 17 00:00:00 2001
From: Jonathan Wren <jonathan@nowandwren.com>
Date: Sat, 24 Sep 2022 15:20:44 -0700
Subject: [PATCH] more cleanup

---
 jrnl/Journal.py                           |  2 +-
 jrnl/encryption/BaseEncryption.py         |  3 ++-
 jrnl/encryption/BasePasswordEncryption.py |  8 ++++---
 jrnl/encryption/Jrnlv1Encryption.py       |  4 ++--
 jrnl/encryption/__init__.py               | 26 +++++++++++++----------
 jrnl/keyring.py                           |  2 +-
 jrnl/prompt.py                            |  2 +-
 7 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/jrnl/Journal.py b/jrnl/Journal.py
index 0e4faeed..4a2dae30 100644
--- a/jrnl/Journal.py
+++ b/jrnl/Journal.py
@@ -81,7 +81,7 @@ class Journal:
 
     def _get_encryption_method(self):
         encryption_method = determine_encryption_method(self.config["encrypt"])
-        self.encryption_method = encryption_method(self.config)
+        self.encryption_method = encryption_method(self.name, self.config)
 
     def _decrypt(self, text):
         if self.encryption_method is None:
diff --git a/jrnl/encryption/BaseEncryption.py b/jrnl/encryption/BaseEncryption.py
index d54440b4..6fcaa2b1 100644
--- a/jrnl/encryption/BaseEncryption.py
+++ b/jrnl/encryption/BaseEncryption.py
@@ -5,7 +5,8 @@ from abc import abstractmethod
 
 
 class BaseEncryption(ABC):
-    def __init__(self, config):
+    def __init__(self, journal_name, config):
+        self._journal_name = journal_name
         self._config = config
 
     @abstractmethod
diff --git a/jrnl/encryption/BasePasswordEncryption.py b/jrnl/encryption/BasePasswordEncryption.py
index 3065fbb7..dc45c3a6 100644
--- a/jrnl/encryption/BasePasswordEncryption.py
+++ b/jrnl/encryption/BasePasswordEncryption.py
@@ -24,8 +24,10 @@ class BasePasswordEncryption(BaseEncryption):
         self._password = None
         self._encoding = "utf-8"
 
-        # Check keyring first to be ready for decryption
-        self._password = get_keyring_password(self._config["journal"])
+        # Check keyring first for password.
+        # That way we'll have it.
+        if keyring_pw := get_keyring_password(self._journal_name):
+            self.password = keyring_pw
 
     @property
     def password(self):
@@ -37,7 +39,7 @@ class BasePasswordEncryption(BaseEncryption):
 
     def encrypt(self, text):
         if self.password is None:
-            self.password = create_password(self._config["journal"])
+            self.password = create_password(self._journal_name)
         return self._encrypt(text)
 
     def decrypt(self, text):
diff --git a/jrnl/encryption/Jrnlv1Encryption.py b/jrnl/encryption/Jrnlv1Encryption.py
index 2e7569db..3ab70b4a 100644
--- a/jrnl/encryption/Jrnlv1Encryption.py
+++ b/jrnl/encryption/Jrnlv1Encryption.py
@@ -15,8 +15,8 @@ class Jrnlv1Encryption(BasePasswordEncryption):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
-    # def _encrypt(self, text: str) -> bytes:
-    # raise NotImplementedError
+    def _encrypt(self, text: str) -> bytes:
+        raise NotImplementedError
 
     def _decrypt(self, text: bytes) -> str | None:
         iv, cipher = text[:16], text[16:]
diff --git a/jrnl/encryption/__init__.py b/jrnl/encryption/__init__.py
index 78f84913..7fbb8e46 100644
--- a/jrnl/encryption/__init__.py
+++ b/jrnl/encryption/__init__.py
@@ -1,18 +1,22 @@
 # Copyright © 2012-2022 jrnl contributors
 # License: https://www.gnu.org/licenses/gpl-3.0.html
-from jrnl.encryption.NoEncryption import NoEncryption
+from importlib import import_module
+
+ENCRYPTION_METHODS = {
+    # config: classname (as string)
+    "default": "NoEncryption",
+    False: "NoEncryption",
+    True: "Jrnlv2Encryption",
+    "jrnlv1": "Jrnlv1Encryption",
+    "jrnlv2": "Jrnlv2Encryption",
+}
 
 
 def determine_encryption_method(config):
-    encryption_method = NoEncryption
-    if config is True or config == "jrnlv2":
-        # Default encryption method
-        from jrnl.encryption.Jrnlv2Encryption import Jrnlv2Encryption
+    key = config
+    if isinstance(config, str):
+        key = config.lower()
 
-        encryption_method = Jrnlv2Encryption
-    elif config == "jrnlv1":
-        from jrnl.encryption.Jrnlv1Encryption import Jrnlv1Encryption
+    my_class = ENCRYPTION_METHODS.get(key, "default")
 
-        encryption_method = Jrnlv1Encryption
-
-    return encryption_method
+    return getattr(import_module(f"jrnl.encryption.{my_class}"), my_class)
diff --git a/jrnl/keyring.py b/jrnl/keyring.py
index 534c117a..250a0bee 100644
--- a/jrnl/keyring.py
+++ b/jrnl/keyring.py
@@ -17,7 +17,7 @@ def get_keyring_password(journal_name: str = "default") -> str | None:
         return None
 
 
-def set_keyring_password(password, journal_name: str = "default"):
+def set_keyring_password(password: str, journal_name: str = "default"):
     try:
         return keyring.set_password("jrnl", journal_name, password)
     except keyring.errors.KeyringError as e:
diff --git a/jrnl/prompt.py b/jrnl/prompt.py
index 7dc701a5..320c5cab 100644
--- a/jrnl/prompt.py
+++ b/jrnl/prompt.py
@@ -37,7 +37,7 @@ def create_password(journal_name: str) -> str:
     if yesno(Message(MsgText.PasswordStoreInKeychain), default=True):
         from jrnl.keyring import set_keyring_password
 
-        set_keyring_password(journal_name, pw)
+        set_keyring_password(pw, journal_name)
 
     return pw