Spaces:
Runtime error
Runtime error
File size: 2,687 Bytes
4eff4ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import os
import json
import logging
from typing import Dict
from cryptography.fernet import Fernet, InvalidToken
logger = logging.getLogger(__name__)
class APIKeyManager:
def __init__(self, data_dir: str):
self.data_dir = data_dir
self.api_keys_file = os.path.join(data_dir, "api_keys.json")
self.key_file = os.path.join(data_dir, ".key")
def get_or_create_key(self) -> bytes:
"""Get or create encryption key for API keys"""
if os.path.exists(self.key_file):
with open(self.key_file, 'rb') as f:
return f.read()
else:
key = Fernet.generate_key()
with open(self.key_file, 'wb') as f:
f.write(key)
return key
def encrypt_api_key(self, api_key: str) -> str:
"""Encrypt an API key"""
if not api_key:
return ""
f = Fernet(self.get_or_create_key())
return f.encrypt(api_key.encode()).decode()
def decrypt_api_key(self, encrypted_key: str) -> str:
"""Decrypt an API key"""
if not encrypted_key:
return ""
f = Fernet(self.get_or_create_key())
return f.decrypt(encrypted_key.encode()).decode()
def save(self, provider: str, api_key: str):
"""Save encrypted API key to file"""
keys = self.load()
keys[provider] = self.encrypt_api_key(api_key)
with open(self.api_keys_file, 'w', encoding="utf-8") as f:
json.dump(keys, f)
def load(self) -> Dict[str, str]:
"""Load and decrypt API keys"""
if not os.path.exists(self.api_keys_file):
return {}
try:
with open(self.api_keys_file, 'r', encoding="utf-8") as f:
encrypted_keys = json.load(f)
except (json.JSONDecodeError, OSError) as e:
# A corrupt/truncated api_keys.json must not crash load() (called on
# startup via app_initializer) — treat it as no stored keys.
logger.warning("Failed to read API keys file: %s", e)
return {}
if not isinstance(encrypted_keys, dict):
# Legacy/wrong shape (e.g. a list) — .items() would raise. Ignore it.
logger.warning("API keys file has unexpected shape (%s); ignoring", type(encrypted_keys).__name__)
return {}
decrypted = {}
for provider, key in encrypted_keys.items():
try:
decrypted[provider] = self.decrypt_api_key(key)
except (InvalidToken, ValueError) as e:
logger.warning("Failed to decrypt API key for %s: %s", provider, e)
return decrypted
|