547 lines
19 KiB
Python
Executable File
547 lines
19 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Secure API Key Vault System for Trax and My-AI-Projects
|
|
Provides encrypted storage, inheritance, and validation of API keys
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import sys
|
|
# Add user site-packages to path for cryptography module
|
|
import site
|
|
sys.path.extend(site.getusersitepackages() if isinstance(site.getusersitepackages(), list) else [site.getusersitepackages()])
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Set, Tuple
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
import hashlib
|
|
import base64
|
|
from getpass import getpass
|
|
import argparse
|
|
import subprocess
|
|
|
|
# Try to import cryptography, provide fallback
|
|
try:
|
|
from cryptography.fernet import Fernet
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2
|
|
from cryptography.hazmat.backends import default_backend
|
|
CRYPTO_AVAILABLE = True
|
|
except ImportError as e:
|
|
CRYPTO_AVAILABLE = False
|
|
# Only print warning when running as main script
|
|
if __name__ == "__main__":
|
|
print(f"⚠️ cryptography not installed. Install with: pip install cryptography")
|
|
|
|
@dataclass
|
|
class KeyMetadata:
|
|
"""Metadata for an API key"""
|
|
name: str
|
|
category: str
|
|
description: str
|
|
required_for: List[str] = field(default_factory=list)
|
|
last_rotated: Optional[str] = None
|
|
expires: Optional[str] = None
|
|
validation_pattern: Optional[str] = None
|
|
|
|
class SecureKeyVault:
|
|
"""
|
|
Secure key vault with encryption and project inheritance
|
|
"""
|
|
|
|
# Standard key definitions for my-ai-projects ecosystem
|
|
STANDARD_KEYS = {
|
|
# AI Models
|
|
"ANTHROPIC_API_KEY": KeyMetadata(
|
|
"ANTHROPIC_API_KEY", "ai", "Claude API",
|
|
["task-master", "main-assistant", "trax-v2"]
|
|
),
|
|
"DEEPSEEK_API_KEY": KeyMetadata(
|
|
"DEEPSEEK_API_KEY", "ai", "DeepSeek API for transcription",
|
|
["youtube-summarizer", "trax-v2"]
|
|
),
|
|
"OPENAI_API_KEY": KeyMetadata(
|
|
"OPENAI_API_KEY", "ai", "OpenAI GPT models",
|
|
["main-assistant"]
|
|
),
|
|
"PERPLEXITY_API_KEY": KeyMetadata(
|
|
"PERPLEXITY_API_KEY", "ai", "Perplexity research API",
|
|
["task-master", "research-agent"]
|
|
),
|
|
"OPENROUTER_API_KEY": KeyMetadata(
|
|
"OPENROUTER_API_KEY", "ai", "OpenRouter multi-model access",
|
|
["main-assistant"]
|
|
),
|
|
"GOOGLE_API_KEY": KeyMetadata(
|
|
"GOOGLE_API_KEY", "ai", "Google Gemini models",
|
|
["youtube-summarizer"]
|
|
),
|
|
"XAI_API_KEY": KeyMetadata(
|
|
"XAI_API_KEY", "ai", "Grok models",
|
|
["task-master"]
|
|
),
|
|
"MISTRAL_API_KEY": KeyMetadata(
|
|
"MISTRAL_API_KEY", "ai", "Mistral models",
|
|
["task-master"]
|
|
),
|
|
|
|
# Services
|
|
"YOUTUBE_API_KEY": KeyMetadata(
|
|
"YOUTUBE_API_KEY", "services", "YouTube Data API",
|
|
["youtube-summarizer", "trax"]
|
|
),
|
|
"SLACK_BOT_TOKEN": KeyMetadata(
|
|
"SLACK_BOT_TOKEN", "services", "Slack bot integration",
|
|
["main-assistant"]
|
|
),
|
|
"GITHUB_TOKEN": KeyMetadata(
|
|
"GITHUB_TOKEN", "services", "GitHub API access",
|
|
["main-assistant"]
|
|
),
|
|
"GITEA_TOKEN": KeyMetadata(
|
|
"GITEA_TOKEN", "services", "Gitea CI/CD",
|
|
["ci-cd"]
|
|
),
|
|
|
|
# Google OAuth
|
|
"GOOGLE_CLIENT_ID": KeyMetadata(
|
|
"GOOGLE_CLIENT_ID", "oauth", "Google OAuth client",
|
|
["main-assistant", "youtube-summarizer"]
|
|
),
|
|
"GOOGLE_CLIENT_SECRET": KeyMetadata(
|
|
"GOOGLE_CLIENT_SECRET", "oauth", "Google OAuth secret",
|
|
["main-assistant", "youtube-summarizer"]
|
|
),
|
|
|
|
# Database
|
|
"DATABASE_URL": KeyMetadata(
|
|
"DATABASE_URL", "database", "PostgreSQL connection string",
|
|
["trax", "main-assistant"]
|
|
),
|
|
"REDIS_URL": KeyMetadata(
|
|
"REDIS_URL", "database", "Redis connection string",
|
|
["cache-layer"]
|
|
),
|
|
}
|
|
|
|
def __init__(self, root_dir: Optional[Path] = None):
|
|
"""Initialize the secure key vault"""
|
|
self.root_dir = root_dir or Path.home() / ".my-ai-keys"
|
|
self.vault_file = self.root_dir / "vault.enc"
|
|
self.metadata_file = self.root_dir / "metadata.json"
|
|
self.master_key_file = self.root_dir / ".master"
|
|
|
|
# Create directories
|
|
self.root_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Initialize encryption
|
|
self.cipher = None
|
|
if CRYPTO_AVAILABLE:
|
|
self._init_encryption()
|
|
|
|
def _init_encryption(self):
|
|
"""Initialize or load encryption key"""
|
|
if self.master_key_file.exists():
|
|
# Load existing key
|
|
with open(self.master_key_file, 'rb') as f:
|
|
key = f.read()
|
|
self.cipher = Fernet(key)
|
|
else:
|
|
# Generate new key with password derivation
|
|
password = getpass("Create vault password: ")
|
|
confirm = getpass("Confirm password: ")
|
|
|
|
if password != confirm:
|
|
print("❌ Passwords don't match")
|
|
sys.exit(1)
|
|
|
|
# Derive key from password
|
|
salt = os.urandom(16)
|
|
kdf = PBKDF2(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend()
|
|
)
|
|
key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
|
|
|
|
# Save key and salt
|
|
with open(self.master_key_file, 'wb') as f:
|
|
f.write(key)
|
|
|
|
# Protect the master key file
|
|
os.chmod(self.master_key_file, 0o600)
|
|
|
|
self.cipher = Fernet(key)
|
|
print("✅ Vault created successfully")
|
|
|
|
def _load_vault(self) -> Dict[str, str]:
|
|
"""Load and decrypt the vault"""
|
|
if not self.vault_file.exists():
|
|
return {}
|
|
|
|
if not self.cipher:
|
|
print("❌ Encryption not available")
|
|
return {}
|
|
|
|
try:
|
|
with open(self.vault_file, 'rb') as f:
|
|
encrypted = f.read()
|
|
|
|
decrypted = self.cipher.decrypt(encrypted)
|
|
return json.loads(decrypted.decode())
|
|
except Exception as e:
|
|
print(f"❌ Failed to decrypt vault: {e}")
|
|
# Try password unlock
|
|
return self._unlock_vault()
|
|
|
|
def _unlock_vault(self) -> Dict[str, str]:
|
|
"""Unlock vault with password"""
|
|
password = getpass("Enter vault password: ")
|
|
|
|
# Re-derive key from password (simplified for demo)
|
|
# In production, store salt separately
|
|
salt = b'my-ai-projects-salt' # Should be stored
|
|
kdf = PBKDF2(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend()
|
|
)
|
|
key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
|
|
|
|
try:
|
|
cipher = Fernet(key)
|
|
with open(self.vault_file, 'rb') as f:
|
|
encrypted = f.read()
|
|
decrypted = cipher.decrypt(encrypted)
|
|
|
|
# Update cipher for future operations
|
|
self.cipher = cipher
|
|
|
|
return json.loads(decrypted.decode())
|
|
except:
|
|
print("❌ Invalid password")
|
|
sys.exit(1)
|
|
|
|
def _save_vault(self, vault: Dict[str, str]):
|
|
"""Encrypt and save the vault"""
|
|
if not self.cipher:
|
|
print("❌ Encryption not available")
|
|
return
|
|
|
|
# Encrypt vault
|
|
data = json.dumps(vault, indent=2)
|
|
encrypted = self.cipher.encrypt(data.encode())
|
|
|
|
# Save encrypted vault
|
|
with open(self.vault_file, 'wb') as f:
|
|
f.write(encrypted)
|
|
|
|
# Protect the vault file
|
|
os.chmod(self.vault_file, 0o600)
|
|
|
|
def add_key(self, name: str, value: str, category: str = None):
|
|
"""Add or update a key in the vault"""
|
|
vault = self._load_vault()
|
|
|
|
# Validate key name
|
|
if name in self.STANDARD_KEYS:
|
|
metadata = self.STANDARD_KEYS[name]
|
|
category = metadata.category
|
|
print(f"📝 Adding standard key: {name} ({metadata.description})")
|
|
else:
|
|
if not category:
|
|
category = "custom"
|
|
print(f"📝 Adding custom key: {name}")
|
|
|
|
# Store key
|
|
vault[name] = value
|
|
self._save_vault(vault)
|
|
|
|
# Update metadata
|
|
self._update_metadata(name, category)
|
|
|
|
print(f"✅ Key '{name}' added to vault")
|
|
|
|
def _update_metadata(self, name: str, category: str):
|
|
"""Update key metadata"""
|
|
metadata = {}
|
|
if self.metadata_file.exists():
|
|
with open(self.metadata_file, 'r') as f:
|
|
metadata = json.load(f)
|
|
|
|
metadata[name] = {
|
|
"category": category,
|
|
"added": datetime.now().isoformat(),
|
|
"last_accessed": None
|
|
}
|
|
|
|
with open(self.metadata_file, 'w') as f:
|
|
json.dump(metadata, f, indent=2)
|
|
|
|
def get_key(self, name: str) -> Optional[str]:
|
|
"""Retrieve a key from the vault"""
|
|
vault = self._load_vault()
|
|
|
|
if name in vault:
|
|
# Update last accessed
|
|
self._update_access_time(name)
|
|
return vault[name]
|
|
|
|
return None
|
|
|
|
def _update_access_time(self, name: str):
|
|
"""Update last access time for a key"""
|
|
if self.metadata_file.exists():
|
|
with open(self.metadata_file, 'r') as f:
|
|
metadata = json.load(f)
|
|
|
|
if name in metadata:
|
|
metadata[name]["last_accessed"] = datetime.now().isoformat()
|
|
|
|
with open(self.metadata_file, 'w') as f:
|
|
json.dump(metadata, f, indent=2)
|
|
|
|
def list_keys(self, category: Optional[str] = None) -> List[str]:
|
|
"""List all keys in the vault"""
|
|
vault = self._load_vault()
|
|
|
|
if not category:
|
|
return list(vault.keys())
|
|
|
|
# Filter by category
|
|
metadata = {}
|
|
if self.metadata_file.exists():
|
|
with open(self.metadata_file, 'r') as f:
|
|
metadata = json.load(f)
|
|
|
|
filtered = []
|
|
for key in vault.keys():
|
|
if key in self.STANDARD_KEYS:
|
|
if self.STANDARD_KEYS[key].category == category:
|
|
filtered.append(key)
|
|
elif key in metadata and metadata[key].get("category") == category:
|
|
filtered.append(key)
|
|
|
|
return filtered
|
|
|
|
def export_to_env(self, output_file: Path, project: Optional[str] = None):
|
|
"""Export keys to .env file format"""
|
|
vault = self._load_vault()
|
|
|
|
# Filter keys for specific project
|
|
keys_to_export = {}
|
|
|
|
if project:
|
|
# Export only keys required for this project
|
|
for key_name, metadata in self.STANDARD_KEYS.items():
|
|
if project in metadata.required_for and key_name in vault:
|
|
keys_to_export[key_name] = vault[key_name]
|
|
else:
|
|
keys_to_export = vault
|
|
|
|
# Write .env file
|
|
with open(output_file, 'w') as f:
|
|
f.write("# API Keys exported from secure vault\n")
|
|
f.write(f"# Generated: {datetime.now().isoformat()}\n")
|
|
if project:
|
|
f.write(f"# Project: {project}\n")
|
|
f.write("\n")
|
|
|
|
# Group by category
|
|
categories = {}
|
|
for key_name, value in keys_to_export.items():
|
|
if key_name in self.STANDARD_KEYS:
|
|
cat = self.STANDARD_KEYS[key_name].category
|
|
else:
|
|
cat = "custom"
|
|
|
|
if cat not in categories:
|
|
categories[cat] = []
|
|
categories[cat].append((key_name, value))
|
|
|
|
# Write grouped keys
|
|
for cat in sorted(categories.keys()):
|
|
f.write(f"# {cat.upper()} KEYS\n")
|
|
for key_name, value in sorted(categories[cat]):
|
|
f.write(f"{key_name}={value}\n")
|
|
f.write("\n")
|
|
|
|
# Protect the .env file
|
|
os.chmod(output_file, 0o600)
|
|
|
|
print(f"✅ Exported {len(keys_to_export)} keys to {output_file}")
|
|
|
|
def import_from_env(self, env_file: Path):
|
|
"""Import keys from .env file"""
|
|
if not env_file.exists():
|
|
print(f"❌ File not found: {env_file}")
|
|
return
|
|
|
|
imported = 0
|
|
with open(env_file, 'r') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#'):
|
|
if '=' in line:
|
|
key, value = line.split('=', 1)
|
|
self.add_key(key, value)
|
|
imported += 1
|
|
|
|
print(f"✅ Imported {imported} keys from {env_file}")
|
|
|
|
def validate_project_keys(self, project: str) -> Tuple[List[str], List[str]]:
|
|
"""Validate that all required keys for a project are present"""
|
|
vault = self._load_vault()
|
|
|
|
required = []
|
|
missing = []
|
|
|
|
for key_name, metadata in self.STANDARD_KEYS.items():
|
|
if project in metadata.required_for:
|
|
required.append(key_name)
|
|
if key_name not in vault:
|
|
missing.append(key_name)
|
|
|
|
return required, missing
|
|
|
|
def rotate_key(self, name: str):
|
|
"""Rotate (regenerate) a key"""
|
|
current = self.get_key(name)
|
|
if not current:
|
|
print(f"❌ Key '{name}' not found")
|
|
return
|
|
|
|
print(f"Current value: {current[:8]}...")
|
|
new_value = input("Enter new value (or press Enter to cancel): ").strip()
|
|
|
|
if new_value:
|
|
self.add_key(name, new_value)
|
|
print(f"✅ Key '{name}' rotated successfully")
|
|
|
|
def sync_to_projects(self, projects: List[str]):
|
|
"""Sync keys to multiple project .env files"""
|
|
workspace_root = Path(__file__).parent.parent.parent.parent
|
|
|
|
for project in projects:
|
|
if project == "root":
|
|
env_file = workspace_root / ".env"
|
|
elif project == "trax":
|
|
env_file = workspace_root / "apps" / "trax" / ".env"
|
|
elif project == "youtube-summarizer":
|
|
env_file = workspace_root / "apps" / "youtube-summarizer" / ".env"
|
|
elif project == "pdf-translator":
|
|
env_file = workspace_root / "pdf-translator" / ".env"
|
|
else:
|
|
print(f"⚠️ Unknown project: {project}")
|
|
continue
|
|
|
|
self.export_to_env(env_file, project)
|
|
print(f"✅ Synced to {project}")
|
|
|
|
def main():
|
|
"""CLI interface for the key vault"""
|
|
parser = argparse.ArgumentParser(description="Secure API Key Vault")
|
|
subparsers = parser.add_subparsers(dest='command', help='Commands')
|
|
|
|
# Add key
|
|
add_parser = subparsers.add_parser('add', help='Add a key to vault')
|
|
add_parser.add_argument('name', help='Key name')
|
|
add_parser.add_argument('--value', help='Key value (will prompt if not provided)')
|
|
add_parser.add_argument('--category', help='Key category')
|
|
|
|
# Get key
|
|
get_parser = subparsers.add_parser('get', help='Get a key from vault')
|
|
get_parser.add_argument('name', help='Key name')
|
|
|
|
# List keys
|
|
list_parser = subparsers.add_parser('list', help='List keys')
|
|
list_parser.add_argument('--category', help='Filter by category')
|
|
|
|
# Import
|
|
import_parser = subparsers.add_parser('import', help='Import from .env file')
|
|
import_parser.add_argument('file', help='Path to .env file')
|
|
|
|
# Export
|
|
export_parser = subparsers.add_parser('export', help='Export to .env file')
|
|
export_parser.add_argument('file', help='Output .env file')
|
|
export_parser.add_argument('--project', help='Filter by project')
|
|
|
|
# Validate
|
|
validate_parser = subparsers.add_parser('validate', help='Validate project keys')
|
|
validate_parser.add_argument('project', help='Project name')
|
|
|
|
# Sync
|
|
sync_parser = subparsers.add_parser('sync', help='Sync to project .env files')
|
|
sync_parser.add_argument('projects', nargs='+', help='Project names')
|
|
|
|
# Rotate
|
|
rotate_parser = subparsers.add_parser('rotate', help='Rotate a key')
|
|
rotate_parser.add_argument('name', help='Key name')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not CRYPTO_AVAILABLE:
|
|
print("❌ cryptography package required")
|
|
print("Install with: pip install cryptography")
|
|
sys.exit(1)
|
|
|
|
vault = SecureKeyVault()
|
|
|
|
if args.command == 'add':
|
|
value = args.value
|
|
if not value:
|
|
value = getpass(f"Enter value for {args.name}: ")
|
|
vault.add_key(args.name, value, args.category)
|
|
|
|
elif args.command == 'get':
|
|
value = vault.get_key(args.name)
|
|
if value:
|
|
print(f"{args.name}={value}")
|
|
else:
|
|
print(f"❌ Key '{args.name}' not found")
|
|
|
|
elif args.command == 'list':
|
|
keys = vault.list_keys(args.category)
|
|
if keys:
|
|
print("\n📋 Keys in vault:")
|
|
for key in sorted(keys):
|
|
if key in vault.STANDARD_KEYS:
|
|
meta = vault.STANDARD_KEYS[key]
|
|
print(f" • {key} ({meta.category}) - {meta.description}")
|
|
else:
|
|
print(f" • {key} (custom)")
|
|
else:
|
|
print("No keys found")
|
|
|
|
elif args.command == 'import':
|
|
vault.import_from_env(Path(args.file))
|
|
|
|
elif args.command == 'export':
|
|
vault.export_to_env(Path(args.file), args.project)
|
|
|
|
elif args.command == 'validate':
|
|
required, missing = vault.validate_project_keys(args.project)
|
|
print(f"\n📋 Project '{args.project}' key validation:")
|
|
print(f" Required: {len(required)} keys")
|
|
if missing:
|
|
print(f" ❌ Missing: {len(missing)} keys")
|
|
for key in missing:
|
|
print(f" • {key}")
|
|
else:
|
|
print(f" ✅ All required keys present")
|
|
|
|
elif args.command == 'sync':
|
|
vault.sync_to_projects(args.projects)
|
|
|
|
elif args.command == 'rotate':
|
|
vault.rotate_key(args.name)
|
|
|
|
else:
|
|
parser.print_help()
|
|
|
|
if __name__ == "__main__":
|
|
main() |