diff --git a/gestalt/vault.py b/gestalt/vault.py index e579e0b..09e6f82 100644 --- a/gestalt/vault.py +++ b/gestalt/vault.py @@ -15,6 +15,7 @@ class Vault(Provider): + def __init__( self, cert: Optional[Tuple[str, str]] = None, @@ -61,9 +62,13 @@ def __do_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() try: self.vault_client.is_authenticated() @@ -74,9 +79,8 @@ def __do_init( if role and 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=role, jwt=jwt) token = self.vault_client.auth.token.lookup_self() if token is not None: kubes_token = ( @@ -87,8 +91,7 @@ def __do_init( 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") @@ -96,13 +99,13 @@ def __do_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() @@ -112,9 +115,12 @@ def stop(self) -> None: def __del__(self) -> None: self.stop() - def get( - self, key: str, path: str, filter: str, sep: Optional[str] = "." - ) -> Union[str, int, float, bool, List[Any]]: + def get(self, + key: str, + path: str, + filter: str, + sep: Optional[str] = "." + ) -> Union[str, int, float, bool, List[Any]]: return retry_call( f=Vault.__do_get, fargs=[self, key, path, filter, sep], @@ -123,9 +129,12 @@ def get( tries=self.tries, ) - def __do_get( - self, key: str, path: str, filter: str, sep: Optional[str] = "." - ) -> Union[str, int, float, bool, List[Any]]: + def __do_get(self, + key: str, + path: str, + filter: str, + sep: Optional[str] = "." + ) -> Union[str, int, float, bool, List[Any]]: """Gets secret from vault Args: key (str): key to get secret from @@ -140,7 +149,8 @@ def __do_get( 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: @@ -157,10 +167,10 @@ def __do_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: @@ -186,13 +196,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 @@ -205,7 +214,8 @@ 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") @@ -217,10 +227,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}")