From 35db8bbfed31e116fc2329b4dfbf78144f66c21a Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Wed, 12 Jun 2024 23:02:31 +0200
Subject: [PATCH] First attempt at actual code

---
 secret_loader/__index__.py | 74 ++++++++++++++++++++++++++++++++++++++
 secret_loader/_version.py  |  1 +
 test/test_init.py          |  6 ++++
 3 files changed, 81 insertions(+)
 create mode 100644 secret_loader/__index__.py
 create mode 100644 secret_loader/_version.py
 create mode 100644 test/test_init.py

diff --git a/secret_loader/__index__.py b/secret_loader/__index__.py
new file mode 100644
index 0000000..1b1d0ef
--- /dev/null
+++ b/secret_loader/__index__.py
@@ -0,0 +1,74 @@
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+try:
+    import hvac
+except ImportError:
+    hvac = None
+
+class SecretLoader:
+    '''
+    Priority order:
+
+    1. Files pointed to by environment variables
+    2. Secrets folder
+    3. Vault instance if configured
+    '''
+
+    def __init__(self, env_key_prefix: str):
+        assert not env_key_prefix.endswith('_')
+        self.env_key_prefix = env_key_prefix
+
+
+        # Setup vault
+        self.vault_client = None
+        if hvac:
+            self.vault_client = hvac.Client(
+                    url=self._load_or_none('VAULT_URL'),
+                    token=self._load_or_none('VAULT_TOKEN'),
+            )
+            self.vault_mount_point=self._load_or_none('VAULT_MOUNT_POINT')
+
+    def load_or_fail(self, env_key: str) -> str:
+        value = self._load_or_none(env_key)
+        if value is None:
+            msg = 'Could not load secret {}'.format(env_key)
+            raise Exception(msg)
+        logger.info('Loaded secret: %s', env_key)
+        return value
+
+    def load(self, env_key: str) -> str | None:
+        value = self._load_or_none(env_key)
+        if value is None:
+            logger.exception('Could not load secret %s', env_key)
+        logger.info('Loaded secret: %s', env_key)
+        return value
+
+    def _load_or_none(self, env_key: str) -> str | None:
+        return self._load_or_none(env_key)
+
+    def _load_or_none_path_or_file(self, env_key: str) -> str | None:
+        # 1. & 2.
+        filepath = os.environ.get(f'{self.env_key_prefix}_{env_key}')
+
+        if filepath is None:
+            filepath = f'./secrets/{env_key.lower()}'
+        try:
+            with open(filepath) as f:
+                value = f.read().strip()
+            return value
+        except Exception:
+            return None
+
+    def _load_or_none_vault(self, env_key: str) -> str | None:
+        if self.vault_client is None:
+            return None
+
+        read_secret_result = self.vault_client.secrets.kv.v1.read_secret(
+            path=env_key.lower(),
+            mount_point=self.vault_mount_point,
+        )
+        return read_secret_result['data']['value']
diff --git a/secret_loader/_version.py b/secret_loader/_version.py
new file mode 100644
index 0000000..b794fd4
--- /dev/null
+++ b/secret_loader/_version.py
@@ -0,0 +1 @@
+__version__ = '0.1.0'
diff --git a/test/test_init.py b/test/test_init.py
new file mode 100644
index 0000000..de2f8f9
--- /dev/null
+++ b/test/test_init.py
@@ -0,0 +1,6 @@
+
+import secret_loader
+
+def test_init():
+    secret_loader.SecretLoader('TEST')
+