Skip to content

Commit

Permalink
yapf formatting.
Browse files Browse the repository at this point in the history
  • Loading branch information
@samson-ichiban
samson-ichiban committed Jan 4, 2024
1 parent 74d200d commit aa77909
Showing 1 changed file with 40 additions and 30 deletions.
70 changes: 40 additions & 30 deletions gestalt/vault.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@


class Vault(Provider):

def __init__(
self,
cert: Optional[Tuple[str, str]] = None,
Expand Down Expand Up @@ -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()
Expand All @@ -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 = (
Expand All @@ -87,22 +91,21 @@ 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")

dynamic_ttl_renew = Thread(
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()

Expand All @@ -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],
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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")
Expand All @@ -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}")

Expand Down

0 comments on commit aa77909

Please sign in to comment.