LaunchLLM / secure_config.py
Bmccloud22's picture
Deploy LaunchLLM - Production AI Training Platform
ec8f374 verified
"""
Secure Configuration - Encrypted API key storage using Fernet encryption
Provides secure storage and retrieval of API keys with encryption at rest.
"""
import os
import json
from pathlib import Path
from typing import Dict, Optional
from cryptography.fernet import Fernet
class SecureConfig:
"""Manages encrypted API key storage"""
def __init__(self, secrets_dir: str = ".secrets"):
self.secrets_dir = Path(secrets_dir)
self.secrets_dir.mkdir(exist_ok=True)
self.key_file = self.secrets_dir / "encryption.key"
self.data_file = self.secrets_dir / "api_keys.enc"
self.fernet = self._load_or_create_key()
def _load_or_create_key(self) -> Fernet:
"""Load existing encryption key or create new one"""
if self.key_file.exists():
with open(self.key_file, "rb") as f:
key = f.read()
else:
key = Fernet.generate_key()
with open(self.key_file, "wb") as f:
f.write(key)
return Fernet(key)
def save_keys(self, api_keys: Dict[str, str]):
"""Save API keys with encryption"""
# Remove empty keys
api_keys = {k: v for k, v in api_keys.items() if v and v.strip()}
# Serialize to JSON
json_data = json.dumps(api_keys)
# Encrypt
encrypted_data = self.fernet.encrypt(json_data.encode())
# Save to file
with open(self.data_file, "wb") as f:
f.write(encrypted_data)
# Also set environment variables for current session
for key, value in api_keys.items():
os.environ[key] = value
def load_keys(self) -> Dict[str, str]:
"""Load and decrypt API keys"""
if not self.data_file.exists():
return {}
try:
# Read encrypted data
with open(self.data_file, "rb") as f:
encrypted_data = f.read()
# Decrypt
decrypted_data = self.fernet.decrypt(encrypted_data)
# Deserialize
api_keys = json.loads(decrypted_data.decode())
# Set environment variables
for key, value in api_keys.items():
os.environ[key] = value
return api_keys
except Exception as e:
print(f"Error loading keys: {e}")
return {}
def get_key(self, key_name: str) -> Optional[str]:
"""Get a specific API key"""
# First try environment variable
value = os.environ.get(key_name)
if value:
return value
# Then try loading from encrypted file
keys = self.load_keys()
return keys.get(key_name)
def get_masked_key(self, key_name: str) -> Optional[str]:
"""Get a masked version of an API key for display"""
key = self.get_key(key_name)
if not key:
return None
# Show first 8 and last 4 characters
if len(key) > 20:
return f"{key[:8]}...{key[-4:]}"
else:
return f"{key[:4]}...{key[-2:]}"
def delete_keys(self):
"""Delete all stored keys"""
if self.data_file.exists():
self.data_file.unlink()
# Also clear environment variables
env_keys = [
"HUGGINGFACE_TOKEN",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"WANDB_API_KEY",
"ELEVEN_API_KEY",
"RUNPOD_API_KEY"
]
for key in env_keys:
os.environ.pop(key, None)
def has_key(self, key_name: str) -> bool:
"""Check if a key exists"""
return self.get_key(key_name) is not None
# Global instance
_config = None
def get_config() -> SecureConfig:
"""Get the global secure config instance"""
global _config
if _config is None:
_config = SecureConfig()
return _config