feat: update vault logic
This commit is contained in:
@@ -2,6 +2,17 @@ from __future__ import annotations
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||||
import hvac
|
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:
|
class VaultClient:
|
||||||
@@ -16,16 +27,38 @@ class VaultClient:
|
|||||||
mount_point: str,
|
mount_point: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._mount_point = mount_point
|
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 = 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]:
|
def read_secret(self,path: str) -> dict[str,Any]:
|
||||||
|
for attempt in range(2):
|
||||||
|
try:
|
||||||
secret = self._client.secrets.kv.v2.read_secret_version(
|
secret = self._client.secrets.kv.v2.read_secret_version(
|
||||||
path=path,
|
path=path,
|
||||||
mount_point=self._mount_point,
|
mount_point=self._mount_point,
|
||||||
)
|
)
|
||||||
return dict(secret.get('data',{}).get('data',{}))
|
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]:
|
def read_many(self,*paths: str) -> dict[str,Any]:
|
||||||
@@ -35,7 +68,7 @@ class VaultClient:
|
|||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
result.update(self.read_secret(path))
|
result.update(self.read_secret(path))
|
||||||
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden):
|
except (hvac.exceptions.InvalidPath,hvac.exceptions.Forbidden,hvac.exceptions.Unauthorized):
|
||||||
continue
|
continue
|
||||||
return result
|
return result
|
||||||
|
|
||||||
@@ -80,9 +113,12 @@ class JwtKeyStore:
|
|||||||
if self._vault_client is None:
|
if self._vault_client is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
current = self._vault_client.read_secret(self._kid_path)
|
current = self._vault_client.read_secret(self._kid_path)
|
||||||
for kid in self._get_configured_kids(current):
|
for kid in self._get_configured_kids(current):
|
||||||
await self.get_public_key_for_kid(kid)
|
await self.get_public_key_for_kid(kid)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(f'JwtKeyStore refresh failed error={exc!s}')
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -112,7 +148,7 @@ class JwtKeyStore:
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
key_data = self._vault_client.read_secret(f'{self._kids_prefix}/{kid}')
|
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
|
return None
|
||||||
|
|
||||||
public_key = (
|
public_key = (
|
||||||
|
|||||||
Reference in New Issue
Block a user