From d0ba5e9c406413d34b183411fea5bb52558cb358 Mon Sep 17 00:00:00 2001 From: Adisun Wheelock <32372583+adisunw@users.noreply.github.com> Date: Wed, 6 Mar 2024 11:11:51 -0500 Subject: [PATCH] Lazily connect to vault (#43) * lazily connect * lazily connect * lazily connect * add to changelog * edit tests * tests * debugging * debugging * linting * no diff * testing * testing * testing * testing * testing * formatting * testing * add unit test * tests * tests * tests * tests * add property * add property * formatting * formatting --- CHANGELOG.md | 5 +++ gestalt/vault.py | 82 ++++++++++++++++++++++++++++--------------- setup.py | 52 ++++++++++++++------------- tests/test_gestalt.py | 11 ++++++ 4 files changed, 96 insertions(+), 54 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de4bc06..d81da77 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) +## [3.4.0] - 2024-03-04 + +### Added +- Lazily authenticates to vault upon first vault request, instead of authentication on object creation. + ## [3.3.6] - 2023-06-11 ### Added diff --git a/gestalt/vault.py b/gestalt/vault.py index 48a4151..3ac79ff 100644 --- a/gestalt/vault.py +++ b/gestalt/vault.py @@ -41,13 +41,31 @@ def __init__( self.dynamic_token_queue: Queue[Tuple[str, str, str]] = Queue() self.kubes_token_queue: Queue[Tuple[str, str, str]] = Queue() - self.vault_client = hvac.Client(url=url, token=token, cert=cert, verify=verify) + self._vault_client: Optional[hvac.Client] = None self._secret_expiry_times: Dict[str, datetime] = dict() - self._secret_values: Dict[str, Union[str, int, float, bool, List[Any]]] = dict() + self._secret_values: Dict[str, Union[str, int, float, bool, + List[Any]]] = dict() + self._is_connected: bool = False + self._role: Optional[str] = role + self._jwt: Optional[str] = jwt + self._url: Optional[str] = url + self._token: Optional[str] = token + self._cert: Optional[Tuple[str, str]] = cert + self._verify: Optional[bool] = verify self.delay = delay self.tries = tries + @property + def vault_client(self) -> hvac.Client: + if self._vault_client is None: + self._vault_client = hvac.Client(url=self._url, + token=self._token, + cert=self._cert, + verify=self._verify) + return self._vault_client + + def connect(self) -> None: try: retry_call( self.vault_client.is_authenticated, @@ -60,11 +78,11 @@ def __init__( "Gestalt Error: Unable to connect to vault with the given configuration" ) - if role and jwt: + if self._role and self._jwt: try: - hvac.api.auth_methods.Kubernetes(self.vault_client.adapter).login( - role=role, jwt=jwt - ) + hvac.api.auth_methods.Kubernetes( + self.vault_client.adapter).login(role=self._role, + jwt=self._jwt) token = retry_call( self.vault_client.auth.token.lookup_self, exceptions=(RuntimeError, Timeout), @@ -75,14 +93,13 @@ def __init__( if token is not None: kubes_token = ( "kubernetes", - token["data"]["id"], # type: ignore + token["data"]["id"], token["data"]["ttl"], - ) # type: ignore + ) self.kubes_token_queue.put(kubes_token) except hvac.exceptions.InvalidPath: raise RuntimeError( - "Gestalt Error: Kubernetes auth couldn't be performed" - ) + "Gestalt Error: Kubernetes auth couldn't be performed") except requests.exceptions.ConnectionError: raise RuntimeError("Gestalt Error: Couldn't connect to Vault") @@ -90,15 +107,16 @@ def __init__( name="dynamic-token-renew", target=self.worker, daemon=True, - args=(self.dynamic_token_queue,), + args=(self.dynamic_token_queue, ), ) # noqa: F841 kubernetes_ttl_renew = Thread( name="kubes-token-renew", target=self.worker, daemon=True, - args=(self.kubes_token_queue,), + args=(self.kubes_token_queue, ), ) kubernetes_ttl_renew.start() + self._is_connected = True def stop(self) -> None: self._run_worker = False @@ -107,7 +125,11 @@ def __del__(self) -> None: self.stop() def get( - self, key: str, path: str, filter: str, sep: Optional[str] = "." + self, + key: str, + path: str, + filter: str, + sep: Optional[str] = "." ) -> Union[str, int, float, bool, List[Any]]: """Gets secret from vault Args: @@ -118,12 +140,15 @@ def get( Returns: secret (str): secret """ + if not self._is_connected: + self.connect() # if the key has been read before and is not a TTL secret if key in self._secret_values and key not in self._secret_expiry_times: return self._secret_values[key] # if the secret can expire but hasn't expired yet - if key in self._secret_expiry_times and not self._is_secret_expired(key): + if key in self._secret_expiry_times and not self._is_secret_expired( + key): return self._secret_values[key] try: @@ -146,10 +171,10 @@ def get( requested_data = response["data"].get("data", response["data"]) except hvac.exceptions.InvalidPath: raise RuntimeError( - "Gestalt Error: The secret path or mount is set incorrectly" - ) + "Gestalt Error: The secret path or mount is set incorrectly") except requests.exceptions.ConnectionError: - raise RuntimeError("Gestalt Error: Gestalt couldn't connect to Vault") + raise RuntimeError( + "Gestalt Error: Gestalt couldn't connect to Vault") except Exception as err: raise RuntimeError(f"Gestalt Error: {err}") if filter is None: @@ -175,13 +200,12 @@ def _is_secret_expired(self, key: str) -> bool: is_expired = now >= secret_expires_dt return is_expired - def _set_secrets_ttl(self, requested_data: Dict[str, Any], key: str) -> None: - last_vault_rotation_str = requested_data["last_vault_rotation"].split(".")[ - 0 - ] # to the nearest second - last_vault_rotation_dt = datetime.strptime( - last_vault_rotation_str, "%Y-%m-%dT%H:%M:%S" - ) + def _set_secrets_ttl(self, requested_data: Dict[str, Any], + key: str) -> None: + last_vault_rotation_str = requested_data["last_vault_rotation"].split( + ".")[0] # to the nearest second + last_vault_rotation_dt = datetime.strptime(last_vault_rotation_str, + "%Y-%m-%dT%H:%M:%S") ttl = requested_data["ttl"] secret_expires_dt = last_vault_rotation_dt + timedelta(seconds=ttl) self._secret_expiry_times[key] = secret_expires_dt @@ -190,11 +214,11 @@ def worker(self, token_queue: Queue) -> None: # type: ignore """ Worker function to renew lease on expiry """ - try: while self._run_worker: if not token_queue.empty(): - token_type, token_id, token_duration = token = token_queue.get() + token_type, token_id, token_duration = token = token_queue.get( + ) if token_type == "kubernetes": self.vault_client.auth.token.renew(token_id) print("kubernetes token for the app has been renewed") @@ -206,10 +230,10 @@ def worker(self, token_queue: Queue) -> None: # type: ignore sleep((token_duration / 3) * 2) except hvac.exceptions.InvalidPath: raise RuntimeError( - "Gestalt Error: The lease path or mount is set incorrectly" - ) + "Gestalt Error: The lease path or mount is set incorrectly") except requests.exceptions.ConnectionError: - raise RuntimeError("Gestalt Error: Gestalt couldn't connect to Vault") + raise RuntimeError( + "Gestalt Error: Gestalt couldn't connect to Vault") except Exception as err: raise RuntimeError(f"Gestalt Error: {err}") diff --git a/setup.py b/setup.py index e30325d..06f42ea 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ def readme(): - with open('README.md') as f: + with open("README.md") as f: return f.read() @@ -10,27 +10,29 @@ def readme(): reqs = filter(lambda x: not x.startswith("-"), reqs_file.readlines()) reqs_list = list(map(lambda x: x.rstrip(), reqs)) -setup(name='gestalt-cfg', - version='3.3.6', - description='A sensible configuration library for Python', - long_description=readme(), - long_description_content_type="text/markdown", - url='https://github.com/clear-street/gestalt', - author='Clear Street', - author_email='engineering@clearstreet.io', - license='MIT', - packages=find_packages(), - python_requires='>=3.6', - install_requires=reqs_list, - classifiers=[ - 'Development Status :: 5 - Production/Stable', - 'Intended Audience :: Developers', - 'License :: OSI Approved :: MIT License', - 'Topic :: Software Development :: Libraries', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Programming Language :: Python :: 3.8', - 'Programming Language :: Python :: 3 :: Only', - 'Operating System :: OS Independent', - ]) +setup( + name="gestalt-cfg", + version="3.4.0", + description="A sensible configuration library for Python", + long_description=readme(), + long_description_content_type="text/markdown", + url="https://github.com/clear-street/gestalt", + author="Clear Street", + author_email="engineering@clearstreet.io", + license="MIT", + packages=find_packages(), + python_requires=">=3.6", + install_requires=reqs_list, + classifiers=[ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Topic :: Software Development :: Libraries", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3 :: Only", + "Operating System :: OS Independent", + ], +) diff --git a/tests/test_gestalt.py b/tests/test_gestalt.py index bbf82cc..b2452b5 100644 --- a/tests/test_gestalt.py +++ b/tests/test_gestalt.py @@ -557,6 +557,15 @@ def test_set_vault_key(nested_setup): assert secret == "ref+vault://secret/data/testnested#.slack.token" +def test_vault_lazy_connect(mock_vault_workers, mock_vault_k8s_auth): + with patch("gestalt.vault.hvac.Client") as mock_client: + v = Vault(role="test-role", jwt="test-jwt") + assert not v._is_connected + v.get("foo", "foo", ".foo") + assert v._is_connected + mock_client().auth.token.lookup_self.assert_called() + + def test_vault_worker_dynamic(mock_vault_workers, mock_vault_k8s_auth): mock_dynamic_renew, mock_k8s_renew = mock_vault_workers @@ -571,6 +580,7 @@ def except_once(self, **kwargs): autospec=True) as mock_sleep: with patch("gestalt.vault.hvac.Client") as mock_client: v = Vault(role="test-role", jwt="test-jwt") + v.connect() mock_k8s_renew.start.assert_called() @@ -602,6 +612,7 @@ def except_once(self, **kwargs): autospec=True) as mock_sleep: with patch("gestalt.vault.hvac.Client") as mock_client: v = Vault(role="test-role", jwt="test-jwt") + v.connect() mock_k8s_renew.start.assert_called()