From dddf0f401f7612bda55d92b28bd3b02f271aaac9 Mon Sep 17 00:00:00 2001
From: Noloquideus <daniilmanukian@gmail.com>
Date: Tue, 12 May 2026 20:08:17 +0300
Subject: [PATCH] feat: update vault logic

---
 src/infrastructure/vault/__init__.py | 58 ++++++++++++++++++++++------
 1 file changed, 47 insertions(+), 11 deletions(-)

diff --git a/src/infrastructure/vault/__init__.py b/src/infrastructure/vault/__init__.py
index 667b2c3..3471d9f 100644
--- a/src/infrastructure/vault/__init__.py
+++ b/src/infrastructure/vault/__init__.py
@@ -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]:
-        secret = self._client.secrets.kv.v2.read_secret_version(
-            path=path,
-            mount_point=self._mount_point,
-        )
-        return dict(secret.get('data',{}).get('data',{}))
+        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
 
-        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)
+        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 = (