feat: update vault logic
This commit is contained in:
@@ -2,6 +2,17 @@ from __future__ import annotations
|
||||
from typing import Any
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
import hvac
|
||||
from src.infrastructure.logger import logger
|
||||
|
||||
|
||||
def _vault_token_renew_failed(exception: Exception) -> bool:
|
||||
if isinstance(exception,(hvac.exceptions.Forbidden,hvac.exceptions.Unauthorized)):
|
||||
return True
|
||||
message = getattr(exception,'message',None) or str(exception)
|
||||
if isinstance(message,str):
|
||||
lower = message.lower()
|
||||
return 'permission denied' in lower or 'invalid token' in lower or '403' in lower
|
||||
return False
|
||||
|
||||
|
||||
class VaultClient:
|
||||
@@ -16,16 +27,38 @@ class VaultClient:
|
||||
mount_point: str,
|
||||
) -> None:
|
||||
self._mount_point = mount_point
|
||||
self._addr = addr
|
||||
self._role_id = role_id
|
||||
self._secret_id = secret_id
|
||||
self._namespace = namespace
|
||||
self._client = hvac.Client(url=addr,namespace=namespace)
|
||||
self._client.auth.approle.login(role_id=role_id,secret_id=secret_id)
|
||||
self._approle_login()
|
||||
|
||||
|
||||
def _approle_login(self) -> None:
|
||||
self._client.auth.approle.login(role_id=self._role_id,secret_id=self._secret_id)
|
||||
|
||||
|
||||
def _renew_or_login(self) -> None:
|
||||
try:
|
||||
self._client.auth.token.renew_self()
|
||||
except Exception:
|
||||
self._approle_login()
|
||||
|
||||
|
||||
def read_secret(self,path: str) -> dict[str,Any]:
|
||||
for attempt in range(2):
|
||||
try:
|
||||
secret = self._client.secrets.kv.v2.read_secret_version(
|
||||
path=path,
|
||||
mount_point=self._mount_point,
|
||||
)
|
||||
return dict(secret.get('data',{}).get('data',{}))
|
||||
except Exception as exc:
|
||||
if attempt == 0 and _vault_token_renew_failed(exc):
|
||||
self._renew_or_login()
|
||||
continue
|
||||
raise
|
||||
|
||||
|
||||
def read_many(self,*paths: str) -> dict[str,Any]:
|
||||
@@ -35,7 +68,7 @@ class VaultClient:
|
||||
continue
|
||||
try:
|
||||
result.update(self.read_secret(path))
|
||||
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden):
|
||||
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden,hvac.exceptions.Unauthorized):
|
||||
continue
|
||||
return result
|
||||
|
||||
@@ -80,9 +113,12 @@ class JwtKeyStore:
|
||||
if self._vault_client is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
current = self._vault_client.read_secret(self._kid_path)
|
||||
for kid in self._get_configured_kids(current):
|
||||
await self.get_public_key_for_kid(kid)
|
||||
except Exception as exc:
|
||||
logger.warning(f'JwtKeyStore refresh failed error={exc!s}')
|
||||
return None
|
||||
|
||||
|
||||
@@ -112,7 +148,7 @@ class JwtKeyStore:
|
||||
|
||||
try:
|
||||
key_data = self._vault_client.read_secret(f'{self._kids_prefix}/{kid}')
|
||||
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden):
|
||||
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden,hvac.exceptions.Unauthorized):
|
||||
return None
|
||||
|
||||
public_key = (
|
||||
|
||||
Reference in New Issue
Block a user