From c74e2c6622774d085da5ac5634bcc1864186a909 Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Fri, 7 Jun 2024 09:01:03 -0700
Subject: [PATCH] Makes it even easier to add new fields in the future

---
 .../management/commands/document_exporter.py  | 14 +++++++----
 .../management/commands/document_importer.py  | 24 ++++++++++++-------
 src/documents/management/commands/mixins.py   | 17 +++++++++++++
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/documents/management/commands/document_exporter.py b/src/documents/management/commands/document_exporter.py
index 6ec565e74..618c1a4e5 100644
--- a/src/documents/management/commands/document_exporter.py
+++ b/src/documents/management/commands/document_exporter.py
@@ -549,13 +549,19 @@ class Command(CryptMixin, BaseCommand):
         """
         Encrypts certain fields in the export.  Currently limited to the mail account password
         """
+
         if self.passphrase:
             self.setup_crypto(passphrase=self.passphrase)
 
-            for mail_account_record in manifest["mail_accounts"]:
-                mail_account_record["fields"]["password"] = self.encrypt_string(
-                    value=mail_account_record["fields"]["password"],
-                )
+            for crypt_config in self.CRYPT_FIELDS:
+                exporter_key = crypt_config["exporter_key"]
+                crypt_fields = crypt_config["fields"]
+                for manifest_record in manifest[exporter_key]:
+                    for field in crypt_fields:
+                        manifest_record["fields"][field] = self.encrypt_string(
+                            value=manifest_record["fields"][field],
+                        )
+
         elif MailAccount.objects.count() > 0:
             self.stdout.write(
                 self.style.NOTICE(
diff --git a/src/documents/management/commands/document_importer.py b/src/documents/management/commands/document_importer.py
index 9aef0b01c..97b73b743 100644
--- a/src/documents/management/commands/document_importer.py
+++ b/src/documents/management/commands/document_importer.py
@@ -404,16 +404,22 @@ class Command(CryptMixin, BaseCommand):
             # Salt has been loaded from metadata.json at this point, so it cannot be None
             self.setup_crypto(passphrase=self.passphrase, salt=self.salt)
 
-            had_an_account = False
+            had_at_least_one_record = False
 
-            for index, record in enumerate(self.manifest):
-                if record["model"] == "paperless_mail.mailaccount":
-                    record["fields"]["password"] = self.decrypt_string(
-                        value=record["fields"]["password"],
-                    )
-                    self.manifest[index] = record
-                    had_an_account = True
-            if had_an_account:
+            for crypt_config in self.CRYPT_FIELDS:
+                importer_model = crypt_config["model_name"]
+                crypt_fields = crypt_config["fields"]
+                for record in filter(
+                    lambda x: x["model"] == importer_model,
+                    self.manifest,
+                ):
+                    had_at_least_one_record = True
+                    for field in crypt_fields:
+                        record["fields"][field] = self.decrypt_string(
+                            value=record["fields"][field],
+                        )
+
+            if had_at_least_one_record:
                 # It's annoying, but the DB is loaded from the JSON directly
                 # Maybe could change that in the future?
                 (self.source / "manifest.json").write_text(
diff --git a/src/documents/management/commands/mixins.py b/src/documents/management/commands/mixins.py
index 1347957be..823631586 100644
--- a/src/documents/management/commands/mixins.py
+++ b/src/documents/management/commands/mixins.py
@@ -2,6 +2,7 @@ import base64
 import os
 from argparse import ArgumentParser
 from typing import Optional
+from typing import TypedDict
 from typing import Union
 
 from cryptography.fernet import Fernet
@@ -16,6 +17,12 @@ from documents.settings import EXPORTER_CRYPTO_SALT_NAME
 from documents.settings import EXPORTER_CRYPTO_SETTINGS_NAME
 
 
+class CryptFields(TypedDict):
+    exporter_key: str
+    model_name: str
+    fields: list[str]
+
+
 class MultiProcessMixin:
     """
     Small class to handle adding an argument and validating it
@@ -86,6 +93,16 @@ class CryptMixin:
     key_size = 32
     kdf_algorithm = "pbkdf2_sha256"
 
+    CRYPT_FIELDS: CryptFields = [
+        {
+            "exporter_key": "mail_accounts",
+            "model_name": "paperless_mail.mailaccount",
+            "fields": [
+                "password",
+            ],
+        },
+    ]
+
     def get_crypt_params(self) -> dict[str, dict[str, Union[str, int]]]:
         return {
             EXPORTER_CRYPTO_SETTINGS_NAME: {