From 2b1d4366d41206c52aef6573a84be2763d9ea62b Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Wed, 4 Sep 2024 19:37:06 +0200
Subject: [PATCH] Configure secrets folder

---
 secret_loader/__init__.py      | 27 +++++++++++++++++++--------
 test/example-secrets/my_secret |  1 +
 test/test_init.py              |  5 +++++
 3 files changed, 25 insertions(+), 8 deletions(-)
 create mode 100644 test/example-secrets/my_secret

diff --git a/secret_loader/__init__.py b/secret_loader/__init__.py
index 4e077df..acb34da 100644
--- a/secret_loader/__init__.py
+++ b/secret_loader/__init__.py
@@ -18,7 +18,9 @@ Secret loading order:
 0. Hardcoded values. **This is purely for debugging, prototyping, and for
    configuring below options.**
 1. Files pointed to by environment variables. Docker friendly.
-2. Secrets folder. Also Docker friendly.
+2. Secrets folder. Also Docker friendly. Defaults to `secrets`, but can be
+   configured through the `SECRETS_DIRECTORY` key (NOTE: passed directly,
+   rather than through a file.)
 3. [Pass: the standard unix password
    manager](https://www.passwordstore.org/). Most suited for personal
    usage; very unsuited for server environments. Requires `pass` installed
@@ -61,14 +63,15 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 ENV_KEY_PREFIX = 'ENV_KEY_PREFIX'
+ENV_KEY_SECRETS_DIRECTORY = 'SECRETS_DIRECTORY'
 
 ENV_KEY_VAULT_URL = 'VAULT_URL'
-ENV_KEY_VAULT_TOKEN = 'VAULT_TOKEN'
+ENV_KEY_VAULT_TOKEN = 'VAULT_TOKEN' #noqa: S105
 ENV_KEY_VAULT_MOUNT_POINT = 'VAULT_MOUNT_POINT'
 
-ENV_KEY_PASS_FOLDER = 'PASS_STORE_SUBFOLDER'
+ENV_KEY_PASS_FOLDER = 'PASS_STORE_SUBFOLDER' #noqa: S105
 
-DEFAULT_SECRETS_DIRECTORY = Path('./secrets/')
+DEFAULT_SECRETS_DIRECTORY = Path('secrets')
 
 ERROR_MESSAGE_FORMAT = """Failed to load secret with key: \033[1m{secret_name}\033[0m
 
@@ -109,6 +112,7 @@ class SecretLoader:
         self.pass_folder = None
         self.vault_client = None
         self.env_key_prefix = None
+        self.secret_folder = None
 
         # Setup environment
         self.env_key_prefix = self._load_or_none(ENV_KEY_PREFIX)
@@ -120,6 +124,9 @@ class SecretLoader:
                 '_',
             ), 'Prefix must not end with _ (this will be added automatically)'
 
+        # Setup secrets path
+        self.secret_folder = Path(self.hardcoded.get(ENV_KEY_SECRETS_DIRECTORY) or self._load_or_none_env(ENV_KEY_SECRETS_DIRECTORY) or DEFAULT_SECRETS_DIRECTORY)
+
         # Setup pass
         self.pass_folder = self._load_or_none(ENV_KEY_PASS_FOLDER)
 
@@ -172,12 +179,13 @@ class SecretLoader:
         Returns `None` if the secret is not present in either the environment
         or the directory.
         """
-        filepath: Path | str | None = os.environ.get(
-            f'{self.env_key_prefix}_{secret_name.upper()}',
-        )
+        filepath: Path | str | None = self._load_or_none_env(secret_name)
 
         if filepath is None:
-            filepath = DEFAULT_SECRETS_DIRECTORY / secret_name.lower()
+            if self.secret_folder is not None:
+                filepath = self.secret_folder / secret_name.lower()
+            else:
+                return None
 
         try:
             with open(filepath) as f:
@@ -185,6 +193,9 @@ class SecretLoader:
         except FileNotFoundError:
             return None
 
+    def _load_or_none_env(self, secret_name) -> str | None:
+        return os.environ.get(f'{self.env_key_prefix}_{secret_name.upper()}')
+
     def _load_or_none_local_password_store(self, secret_name: str) -> str | None:
         """Load secret from the `pass` password manager.
 
diff --git a/test/example-secrets/my_secret b/test/example-secrets/my_secret
new file mode 100644
index 0000000..f793152
--- /dev/null
+++ b/test/example-secrets/my_secret
@@ -0,0 +1 @@
+HELLO SECRET
diff --git a/test/test_init.py b/test/test_init.py
index 37f50b4..26adc23 100644
--- a/test/test_init.py
+++ b/test/test_init.py
@@ -7,6 +7,11 @@ def test_hardcoded():
     assert loader.load('KEY') == 'VALUE'
 
 
+def test_lookup_secrets_dir():
+    loader = secret_loader.SecretLoader(SECRETS_DIRECTORY='test/example-secrets')
+    assert loader.load('MY_SECRET') == 'HELLO SECRET'
+
+
 def test_lookup_unknown():
     loader = secret_loader.SecretLoader()
     assert loader.load('UNKNOWN') is None