manage_api_keys.py

#!/usr/bin/env python3
"""
Simple API Key Management Script for Agentic AI System.
Provides easy commands to add, remove, list, and manage API keys.
"""

import json
import sys
import string
from pathlib import Path
from typing import Dict, Any

CONFIG_FILE = Path("config/api_keys.json")

def validate_api_key(key: str) -> bool:
    """Validate API key format."""
    if not key or len(key) < 8:
        return False
    # Allow alphanumeric, hyphens, underscores
    allowed_chars = string.ascii_letters + string.digits + "-_"
    return all(c in allowed_chars for c in key)

def load_config() -> Dict[str, Any]:
    """Load current API key configuration."""
    if CONFIG_FILE.exists():
        with open(CONFIG_FILE, 'r') as f:
            return json.load(f)
    else:
        return {
            "api_keys": {},
            "settings": {
                "default_rate_limit": 100,
                "token_expiry_days": 365,
                "require_https": False
            }
        }

def save_config(config: Dict[str, Any]) -> None:
    """Save API key configuration."""
    CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(CONFIG_FILE, 'w') as f:
        json.dump(config, f, indent=2)
    print(f"✅ Configuration saved to {CONFIG_FILE}")

def list_keys() -> None:
    """List all API keys."""
    config = load_config()
    api_keys = config.get("api_keys", {})
    
    if not api_keys:
        print("❌ No API keys found")
        return
    
    print("🔑 Current API Keys:")
    print("-" * 60)
    
    for key, info in api_keys.items():
        status = "✅ Enabled" if info.get("enabled", True) else "❌ Disabled"
        masked_key = key[:8] + "..." + key[-4:] if len(key) > 12 else key
        
        print(f"Key: {masked_key}")
        print(f"  Name: {info.get('name', 'Unknown')}")
        print(f"  Permissions: {', '.join(info.get('permissions', []))}")
        print(f"  Rate Limit: {info.get('rate_limit', 100)}/hour")
        print(f"  Status: {status}")
        print(f"  Description: {info.get('description', 'No description')}")
        print()

def add_key(api_key: str, name: str, permissions: list = None, rate_limit: int = 100, description: str = "") -> None:
    """Add a new API key."""
    config = load_config()
    
    # Validate the API key
    if not validate_api_key(api_key):
        print("❌ Invalid API key format. Must be at least 8 characters and contain only letters, numbers, hyphens, and underscores.")
        return
    
    # Check if key already exists
    if api_key in config.get("api_keys", {}):
        print(f"❌ API key already exists: {api_key}")
        return
    
    # Set default permissions
    if permissions is None:
        permissions = ["read", "write"]
    
    # Add key to config
    config["api_keys"][api_key] = {
        "name": name,
        "permissions": permissions,
        "description": description,
        "rate_limit": rate_limit,
        "enabled": True
    }
    
    save_config(config)
    
    print("🎉 API key added successfully!")
    print(f"Key: {api_key}")
    print(f"Name: {name}")
    print(f"Permissions: {', '.join(permissions)}")
    print(f"Rate Limit: {rate_limit}/hour")

def remove_key(key_pattern: str) -> None:
    """Remove an API key by pattern matching."""
    config = load_config()
    api_keys = config.get("api_keys", {})
    
    # Find matching keys
    matches = [key for key in api_keys.keys() if key_pattern in key]
    
    if not matches:
        print(f"❌ No keys found matching: {key_pattern}")
        return
    
    if len(matches) > 1:
        print(f"❌ Multiple keys match '{key_pattern}'. Be more specific:")
        for key in matches:
            masked_key = key[:8] + "..." + key[-4:]
            print(f"  {masked_key} ({api_keys[key].get('name', 'Unknown')})")
        return
    
    # Remove the key
    key_to_remove = matches[0]
    removed_info = api_keys.pop(key_to_remove)
    config["api_keys"] = api_keys
    
    save_config(config)
    
    print(f"✅ Removed API key: {removed_info.get('name', 'Unknown')}")

def disable_key(key_pattern: str) -> None:
    """Disable an API key."""
    config = load_config()
    api_keys = config.get("api_keys", {})
    
    # Find matching keys
    matches = [key for key in api_keys.keys() if key_pattern in key]
    
    if not matches:
        print(f"❌ No keys found matching: {key_pattern}")
        return
    
    if len(matches) > 1:
        print(f"❌ Multiple keys match '{key_pattern}'. Be more specific.")
        return
    
    # Disable the key
    key_to_disable = matches[0]
    api_keys[key_to_disable]["enabled"] = False
    
    save_config(config)
    
    print(f"✅ Disabled API key: {api_keys[key_to_disable].get('name', 'Unknown')}")

def enable_key(key_pattern: str) -> None:
    """Enable an API key."""
    config = load_config()
    api_keys = config.get("api_keys", {})
    
    # Find matching keys
    matches = [key for key in api_keys.keys() if key_pattern in key]
    
    if not matches:
        print(f"❌ No keys found matching: {key_pattern}")
        return
    
    if len(matches) > 1:
        print(f"❌ Multiple keys match '{key_pattern}'. Be more specific.")
        return
    
    # Enable the key
    key_to_enable = matches[0]
    api_keys[key_to_enable]["enabled"] = True
    
    save_config(config)
    
    print(f"✅ Enabled API key: {api_keys[key_to_enable].get('name', 'Unknown')}")

def create_default_config() -> None:
    """Create default configuration with demo keys."""
    config = {
        "api_keys": {
            "demo-key-12345": {
                "name": "Demo User",
                "permissions": ["read", "write"],
                "description": "Default demo key for testing and development",
                "rate_limit": 100,
                "enabled": True
            },
            "admin-key-67890": {
                "name": "Admin User",
                "permissions": ["read", "write", "admin"],
                "description": "Admin key with full system access",
                "rate_limit": 1000,
                "enabled": True
            }
        },
        "settings": {
            "default_rate_limit": 100,
            "token_expiry_days": 365,
            "require_https": False
        }
    }
    
    save_config(config)
    print("✅ Created default API key configuration")

def show_help() -> None:
    """Show help information."""
    print("🔑 API Key Management for Agentic AI System")
    print("=" * 50)
    print()
    print("Commands:")
    print("  list                              - List all API keys")
    print("  add <api_key> <name>             - Add new API key with default permissions")
    print("  add-admin <api_key> <name>       - Add new admin API key")
    print("  remove <key_pattern>             - Remove API key (partial match)")
    print("  disable <key_pattern>            - Disable API key")
    print("  enable <key_pattern>             - Enable API key")
    print("  init                             - Create default configuration")
    print("  help                             - Show this help")
    print()
    print("Examples:")
    print("  python manage_api_keys.py list")
    print("  python manage_api_keys.py add 'prod-api-key-2024' 'Production Client'")
    print("  python manage_api_keys.py add-admin 'admin-key-secure-2024' 'System Admin'")
    print("  python manage_api_keys.py remove demo-key")
    print("  python manage_api_keys.py disable demo-key-12345")
    print()
    print("Notes:")
    print("  - API keys must be at least 8 characters")
    print("  - Use letters, numbers, hyphens, and underscores only")
    print("  - Choose secure, unique keys for production")
    print()
    print("Configuration file: config/api_keys.json")

def main():
    """Main function."""
    if len(sys.argv) < 2:
        show_help()
        return
    
    command = sys.argv[1].lower()
    
    try:
        if command == "list":
            list_keys()
        
        elif command == "add":
            if len(sys.argv) < 4:
                print("❌ Usage: add <api_key> <name>")
                return
            api_key = sys.argv[2]
            name = sys.argv[3]
            add_key(api_key, name)
        
        elif command == "add-admin":
            if len(sys.argv) < 4:
                print("❌ Usage: add-admin <api_key> <name>")
                return
            api_key = sys.argv[2]
            name = sys.argv[3]
            add_key(api_key, name, permissions=["read", "write", "admin"], rate_limit=1000, description="Admin user with full access")
        
        elif command == "remove":
            if len(sys.argv) < 3:
                print("❌ Usage: remove <key_pattern>")
                return
            pattern = sys.argv[2]
            remove_key(pattern)
        
        elif command == "disable":
            if len(sys.argv) < 3:
                print("❌ Usage: disable <key_pattern>")
                return
            pattern = sys.argv[2]
            disable_key(pattern)
        
        elif command == "enable":
            if len(sys.argv) < 3:
                print("❌ Usage: enable <key_pattern>")
                return
            pattern = sys.argv[2]
            enable_key(pattern)
        
        elif command == "init":
            create_default_config()
        
        elif command == "help":
            show_help()
        
        else:
            print(f"❌ Unknown command: {command}")
            show_help()
    
    except Exception as e:
        print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()