From 689f9497e31bb0288c3c9f36a13a05f92383ff7a Mon Sep 17 00:00:00 2001 From: Adisun Wheelock Date: Mon, 4 Mar 2024 15:45:40 +0000 Subject: [PATCH] lazy connect --- gestalt/vault.py | 92 ++++++++++++++++++++++++++---------------------- 1 file changed, 49 insertions(+), 43 deletions(-) diff --git a/gestalt/vault.py b/gestalt/vault.py index 6c0a76b..3bbd898 100644 --- a/gestalt/vault.py +++ b/gestalt/vault.py @@ -41,17 +41,16 @@ 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 = hvac.Client(url=url, token=token, cert=cert, verify=verify) 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.delay = delay self.tries = tries + def connect(self) -> None: try: retry_call( self.vault_client.is_authenticated, @@ -64,10 +63,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), @@ -78,25 +78,31 @@ def __init__( if token is not None: kubes_token = ( "kubernetes", - token['data']['id'], # type: ignore - token['data']['ttl']) # type: ignore + token["data"]["id"], # type: ignore + 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") dynamic_ttl_renew = Thread( - name='dynamic-token-renew', + name="dynamic-token-renew", target=self.worker, daemon=True, - 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.dynamic_token_queue,), + ) # noqa: F841 + kubernetes_ttl_renew = Thread( + name="kubes-token-renew", + target=self.worker, + daemon=True, + args=(self.kubes_token_queue,), + ) kubernetes_ttl_renew.start() + self._is_connected = True def stop(self) -> None: self._run_worker = False @@ -105,11 +111,7 @@ 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: @@ -120,13 +122,14 @@ 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: @@ -139,17 +142,20 @@ def get( ) if response is None: raise RuntimeError("Gestalt Error: No secrets found") - if response['lease_id']: - dynamic_token = ("dynamic", response['lease_id'], - response['lease_duration']) + if response["lease_id"]: + dynamic_token = ( + "dynamic", + response["lease_id"], + response["lease_duration"], + ) self.dynamic_token_queue.put_nowait(dynamic_token) 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,12 +181,13 @@ 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 @@ -193,8 +200,7 @@ def worker(self, token_queue: Queue) -> None: # type: ignore 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 +212,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}")