#!/usr/bin/env python3 """ API Key Manager - A production-ready module for generating, validating, and rate-limiting API keys. Features: - Generate secure API keys - Validate API keys - Rate limiting with configurable window - Usage tracking with SQLite backend - Command-line interface Usage: python api_key_manager.py generate [--description DESC] python api_key_manager.py validate KEY python api_key_manager.py status KEY [--minutes MINUTES] python api_key_manager.py delete KEY """ import argparse import sqlite3 import sys import time import uuid from datetime import datetime, timedelta from typing import Optional, Tuple # Database configuration DB_NAME = "api_keys.db" RATE_LIMIT_WINDOW = 60 # Rate limit window in seconds MAX_REQUESTS = 100 # Max requests per window class APIKeyManager: """Manages API keys with SQLite backend.""" def __init__(self, db_path: str = DB_NAME): """Initialize the API key manager with database connection. Args: db_path: Path to the SQLite database file """ self.db_path = db_path self.conn = sqlite3.connect(self.db_path, check_same_thread=False) self.conn.row_factory = sqlite3.Row self._initialize_database() def _initialize_database(self) -> None: """Create necessary database tables if they don't exist.""" with self.conn: self.conn.execute( """ CREATE TABLE IF NOT EXISTS api_keys ( key TEXT PRIMARY KEY, description TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, last_used TIMESTAMP, is_active INTEGER DEFAULT 1 ) """ ) self.conn.execute( """ CREATE TABLE IF NOT EXISTS usage_logs ( id INTEGER PRIMARY KEY AUTOINCREMENT, api_key TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (api_key) REFERENCES api_keys (key) ) """ ) def generate_key(self, description: Optional[str] = None) -> str: """Generate a new API key. Args: description: Optional description for the key Returns: Newly generated API key """ api_key = f"ak_{uuid.uuid4().hex[:32]}" try: with self.conn: self.conn.execute( "INSERT INTO api_keys (key, description) VALUES (?, ?)", (api_key, description), ) return api_key except sqlite3.IntegrityError as e: raise ValueError("Failed to generate unique API key") from e def validate_key(self, api_key: str) -> bool: """Validate an API key. Args: api_key: API key to validate Returns: True if valid, False otherwise """ cursor = self.conn.execute( "SELECT is_active FROM api_keys WHERE key = ?", (api_key,) ) result = cursor.fetchone() return bool(result and result["is_active"]) def log_usage(self, api_key: str) -> None: """Log API key usage and update last used timestamp. Args: api_key: API key that was used """ try: with self.conn: self.conn.execute( "UPDATE api_keys SET last_used = CURRENT_TIMESTAMP WHERE key = ?", (api_key,), ) self.conn.execute( "INSERT INTO usage_logs (api_key) VALUES (?)", (api_key,), ) except sqlite3.Error as e: raise RuntimeError(f"Failed to log usage: {e}") from e def check_rate_limit(self, api_key: str) -> Tuple[bool, int]: """Check if API key has exceeded rate limit. Args: api_key: API key to check Returns: A tuple containing a boolean indicating if the rate limit is exceeded, and the number of requests in the current window. """ now = datetime.now() window_start = now - timedelta(seconds=RATE_LIMIT_WINDOW) cursor = self.conn.execute( """ SELECT COUNT(*) FROM usage_logs WHERE api_key = ? AND timestamp >= ? """, (api_key, window_start), ) result = cursor.fetchone() request_count = result[0] if result else 0 return request_count > MAX_REQUESTS, request_count def get_status(self, api_key: str, minutes: int = 0) -> Tuple[Optional[datetime], int, bool]: """Get the status of an API key. Args: api_key: The API key to check. minutes: The number of minutes to check the rate limit window. Returns: A tuple containing: - The last used timestamp (datetime object), or None if never used. - The number of requests in the current window. - True if the rate limit is exceeded, False otherwise. """ cursor = self.conn.execute( "SELECT last_used FROM api_keys WHERE key = ?", (api_key,) ) last_used_row = cursor.fetchone() last_used = last_used_row["last_used"] if last_used_row else None exceeded, request_count = self.check_rate_limit(api_key) return last_used, request_count, exceeded def delete_key(self, api_key: str) -> bool: """Delete an API key. Args: api_key: The API key to delete. Returns: True if the key was deleted, False otherwise. """ try: with self.conn: self.conn.execute( "DELETE FROM api_keys WHERE key = ?", (api_key,), ) return self.conn.total_changes > 0 except sqlite3.Error: return False def main(): parser = argparse.ArgumentParser( description="API Key Manager - Generate, validate, and manage API keys." ) subparsers = parser.add_subparsers(dest="command", help="Available commands") # Generate command generate_parser = subparsers.add_parser( "generate", help="Generate a new API key" ) generate_parser.add_argument( "--description", "-d", help="Description for the API key", required=False ) # Validate command validate_parser = subparsers.add_parser("validate", help="Validate an API key") validate_parser.add_argument("key", help="The API key to validate") # Status command status_parser = subparsers.add_parser("status", help="Get the status of an API key") status_parser.add_argument("key", help="The API key to check") status_parser.add_argument( "--minutes", "-m", type=int, default=0, help="Check rate limit within the last X minutes", ) # Delete command delete_parser = subparsers.add_parser("delete", help="Delete an API key") delete_parser.add_argument("key", help="The API key to delete") args = parser.parse_args() api_key_manager = APIKeyManager() if args.command == "generate": try: key = api_key_manager.generate_key(args.description) print(f"Generated API key: {key}") except ValueError as e: print(f"Error: {e}") sys.exit(1) elif args.command == "validate": is_valid = api_key_manager.validate_key(args.key) if is_valid: print("API key is valid") else: print("API key is invalid") elif args.command == "status": last_used, request_count, exceeded = api_key_manager.get_status( args.key, args.minutes ) print(f"Key: {args.key}") print(f"Requests in the last {RATE_LIMIT_WINDOW} seconds: {request_count}") if last_used: print(f"Last used: {last_used}") else: print("Never used") if exceeded: print("Rate limit exceeded") else: print("Rate limit not exceeded") elif args.command == "delete": if api_key_manager.delete_key(args.key): print("API key deleted") else: print("Failed to delete API key or key not found") else: parser.print_help() sys.exit(1) if __name__ == "__main__": main()