add logic to decrypt .env.vault. load_dotenv signature is the same as python-dotenv

Author: nsnguyen

.gitignore

@@ -104,6 +104,8 @@ celerybeat.pid
 # Environments
 .env
 .venv
+.env.vault
+.env.me
 env/
 venv/
 ENV/
@@ -127,3 +129,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.vscode/

Makefile

@@ -15,6 +15,14 @@ clean-pyc:
 build: clean
 	python setup.py sdist bdist_wheel
 
+uninstall_local:
+	pip uninstall python-dotenv-vault -y
+
+install_local:
+	pip install .
+
+test: uninstall_local build install_local
+
 release: build
 	twine check dist/*
 	twine upload dist/*

README.md

@@ -8,3 +8,9 @@ dotenv-vault-python allows your app to sync the `.env` file to the cloud. This m
 ```shell
 pip install python-dotenv-vault --no-cache-dir
 ```
+
+```shell
+from dotenv_vault import load_dotenv
+
+load_dotenv()
+```

setup.py

@@ -4,7 +4,7 @@
 
 src = {}
 dir = os.path.abspath(os.path.dirname(__file__))
-with open(os.path.join(dir, "src/dotenv", "__version__.py"), "r") as f:
+with open(os.path.join(dir, "src/dotenv_vault", "__version__.py"), "r") as f:
     exec(f.read(), src)
 
 def read_files(files):
@@ -39,5 +39,8 @@ def read_files(files):
     'python',
     'dotenv-vault'
     ],
-    install_requires=[],
+    install_requires=[
+        'python-dotenv~=0.21.0',
+        'cryptography~=38.0.1'
+    ],
 )

src/dotenv/__init__.py

@@ -0,0 +1 @@
+from .main import load_dotenv

src/dotenv/main.py

@@ -0,0 +1,33 @@
+import os
+import logging
+from typing import (IO, Optional,Union)
+from dotenv.main import load_dotenv as load_dotenv_file
+
+from .vault import DotEnvVault
+
+logging.basicConfig(level = logging.INFO)
+
+logger = logging.getLogger(__name__)
+
+
+def load_dotenv(
+    dotenv_path: Union[str, os.PathLike, None] = None,
+    stream: Optional[IO[str]] = None,
+    verbose: bool = False,
+    override: bool = False,
+    interpolate: bool = True,
+    encoding: Optional[str] = "utf-8",
+) -> bool:
+    """
+    parameters are the same as python-dotenv library.
+    This is to inject the parameters to evironment variables.
+    """
+    dotenv_vault = DotEnvVault()
+    logger.info(f'dotenv_key:{dotenv_vault.dotenv_key}')
+    if dotenv_vault.dotenv_key:
+        logger.info('Getting .env from vault.')
+        vault_stream = dotenv_vault.parsed_vault()
+        return load_dotenv_file(stream=vault_stream)
+    else:
+        logger.info('Getting .env from local.')
+        return load_dotenv_file()

src/dotenv_vault/__init__.py

@@ -0,0 +1,84 @@
+
+import logging
+import os
+import sys
+import io
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from base64 import b64decode
+from urllib.parse import urlparse, parse_qsl
+
+from dotenv.main import DotEnv, find_dotenv, load_dotenv
+
+
+logging.basicConfig(level = logging.INFO)
+
+logger = logging.getLogger(__name__)
+
+
+class DotEnvVaultError(Exception):
+    pass
+
+
+class DotEnvVault(): #vault stuff
+    def __init__(self) -> None:
+        logger.info('initializing DotEnvVault')
+        self.dotenv_key = os.environ.get('DOTENV_KEY')
+        
+
+    def parsed_vault(self) -> bytes:
+        """
+        Parse information from DOTENV_KEY, and decrypt vault key.
+        """
+        if self.dotenv_key is None: raise DotEnvVaultError("NOT_FOUND_DOTENV_KEY: Cannot find ENV['DOTENV_KEY']")
+            
+        # .env.vault needs to be present.
+        env_vault_path = find_dotenv(filename='.env.vault')
+        if env_vault_path == '':
+            raise DotEnvVaultError("ENV_VAULT_NOT_FOUND: .env.vault is not present.")
+
+        # parse DOTENV_KEY, format is a URI
+        uri = urlparse(self.dotenv_key)
+        # Get encrypted key
+        key = uri.password
+        # Get environment from query params.
+        params = dict(parse_qsl(uri.query))
+        vault_environment = params.get('environment').upper()
+
+        if vault_environment is None or vault_environment not in ['PRODUCTION', 'DEVELOPMENT', 'CI', 'STAGING']:
+            raise DotEnvVaultError('Incorrect Vault Environment.')
+
+        # Getting ciphertext from correct environment in .env.vault
+        environment_key = f'DOTENV_VAULT_{vault_environment}'
+        logging.info(f'Getting {environment_key}.')
+
+        # use python-dotenv library class.
+        dotenv = DotEnv(dotenv_path=env_vault_path)
+        ciphertext = dotenv.dict().get(environment_key)
+
+        decrypted = self._decrypt(ciphertext=ciphertext, key=key)
+        return self._to_text_stream(decrypted)
+
+
+    def _decrypt(self, ciphertext: str, key: str) -> bytes:
+        """
+        decrypt method will decrypt via AES-GCM 
+        return: decrypted keys in bytes
+        """
+        _key = key[4:]
+        if len(_key) < 64: raise DotEnvVault('INVALID_DOTENV_KEY: Key part must be 64 characters long (or more)')
+
+        _key = bytes.fromhex(_key)
+        ciphertext = b64decode(ciphertext)
+
+        aesgcm = AESGCM(_key)
+        return aesgcm.decrypt(ciphertext[:12], ciphertext[12:], b'')
+
+    def _to_text_stream(self, decrypted_obj: bytes) -> io.StringIO:
+        """
+        convert decrypted object (in bytes) to io.StringIO format.
+        Python-dotenv is expecting stream to be text stream (such as `io.StringIO`).
+        return: io.StringIO
+        """
+        decoded_str = decrypted_obj.decode('utf-8')
+        return io.StringIO(decoded_str)