src/pqc_agent_wallet/integrations/env_shim.py
| 1 | """Install an os.environ shim that falls back to the wallet for missing keys. |
| 2 | |
| 3 | Usage: |
| 4 | from pqc_agent_wallet.integrations.env_shim import install_env_shim |
| 5 | install_env_shim(wallet) |
| 6 | |
| 7 | # Now any code that does os.getenv('OPENAI_API_KEY') will fall back |
| 8 | # to wallet.get('openai_api_key') if it isn't in os.environ. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import os |
| 14 | |
| 15 | from pqc_agent_wallet.errors import CredentialNotFoundError, WalletLockedError |
| 16 | from pqc_agent_wallet.vault import Wallet |
| 17 | |
| 18 | _original_getenv = None |
| 19 | |
| 20 | |
| 21 | def install_env_shim(wallet: Wallet) -> None: |
| 22 | """Monkey-patch os.getenv to fall back to the wallet.""" |
| 23 | global _original_getenv |
| 24 | if _original_getenv is not None: |
| 25 | return # already installed |
| 26 | _original_getenv = os.getenv |
| 27 | |
| 28 | def shim(key: str, default: str | None = None) -> str | None: |
| 29 | val = _original_getenv(key, None) if _original_getenv else None |
| 30 | if val is not None: |
| 31 | return val |
| 32 | if not wallet.is_unlocked: |
| 33 | return default |
| 34 | try: |
| 35 | return wallet.get(key.lower()) |
| 36 | except (CredentialNotFoundError, WalletLockedError): |
| 37 | return default |
| 38 | |
| 39 | os.getenv = shim # type: ignore[assignment] |
| 40 | |
| 41 | |
| 42 | def uninstall_env_shim() -> None: |
| 43 | """Restore the original os.getenv.""" |
| 44 | global _original_getenv |
| 45 | if _original_getenv is None: |
| 46 | return |
| 47 | os.getenv = _original_getenv # type: ignore[assignment] |
| 48 | _original_getenv = None |
| 49 | |