api-key-manager 2.1.0__py3-none-any.whl

key_manager/cli.py

@@ -0,0 +1,270 @@
+import argparse
+import asyncio
+import json
+import sys
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from rich.console import Console
+from rich.table import Table
+
+from key_manager.config import load_config
+from key_manager.parser import import_keys
+from key_manager.validator import validate_keys
+from key_manager.checker import run_check
+from key_manager.tester import run_test
+from key_manager.proxy import get_proxy
+from key_manager.storage import KeyStore
+from key_manager.errors import KeyManagerError, ErrorCode
+
+console = Console()
+
+
+def _load_keys(config: dict) -> dict:
+    """Load keys data using KeyStore, falling back to direct JSON read."""
+    try:
+        return KeyStore(config["storage"]["keys_file"], config).load()
+    except Exception:
+        keys_path = Path(config["storage"]["keys_file"])
+        if not keys_path.exists():
+            return {"keys": {}}
+        with open(keys_path, "r", encoding="utf-8") as f:
+            return json.load(f)
+
+
+def cmd_import(args, config):
+    new, dupes, errors = import_keys(
+        file_path=args.file,
+        directory=args.dir or config["scan"]["directories"][0],
+        batch=args.batch,
+        keys_file=config["storage"]["keys_file"]
+    )
+    console.print(f"[green]Import complete:[/green] {new} new, {dupes} duplicates")
+    for err in errors:
+        console.print(f"[red]Error:[/red] {err}")
+
+
+def cmd_check(args, config):
+    proxy = get_proxy(config.get("proxy", ""))
+    if proxy:
+        console.print(f"[blue]Using proxy:[/blue] {proxy}")
+    results = asyncio.run(run_check(
+        keys_file=config["storage"]["keys_file"],
+        results_file=config["storage"]["check_results_file"],
+        logs_dir=config["storage"]["logs_dir"],
+        concurrency=config["check"]["concurrency"],
+        timeout=config["check"]["timeout_seconds"],
+        proxy=proxy or None,
+        retry_failed=config["check"]["retry_failed"],
+        retry_count=config["check"]["retry_count"]
+    ))
+    console.print(f"\n[bold]Check Results[/bold]")
+    console.print(f"Total: {results['total']}")
+    console.print(f"Valid: {results['summary']['valid']['count']}")
+    console.print(f"Invalid: {results['summary']['invalid']['count']}")
+    console.print(f"Error: {results['summary']['error']['count']}")
+
+
+def cmd_test(args, config):
+    proxy = get_proxy(config.get("proxy", ""))
+    if proxy:
+        console.print(f"[blue]Using proxy:[/blue] {proxy}")
+    results = asyncio.run(run_test(
+        keys_file=config["storage"]["keys_file"],
+        results_file=config["storage"]["test_results_file"],
+        logs_dir=config["storage"]["logs_dir"],
+        timeout=config["test"]["concurrency_timeout_seconds"],
+        proxy=proxy or None,
+        token_test=not args.skip_token,
+        concurrency_test=not args.skip_concurrency,
+        token_steps=config["test"]["token_steps"],
+        concurrency_steps=config["test"]["concurrency_steps"],
+        provider_filter=args.provider,
+        single_key=args.key
+    ))
+    console.print(f"\n[bold]Test Results[/bold]")
+    console.print(f"Total tested: {results['total_tested']}")
+
+
+def cmd_list(args, config):
+    keys_path = Path(config["storage"]["keys_file"])
+    if not keys_path.exists():
+        console.print("[yellow]No keys file found. Run 'import' first.[/yellow]")
+        return
+
+    try:
+        data = _load_keys(config)
+    except KeyManagerError as e:
+        console.print(f"[red]Error:[/red] {e.message}")
+        return
+
+    table = Table(title="API Keys")
+    table.add_column("Key", style="cyan")
+    table.add_column("Provider", style="green")
+    table.add_column("Status", style="yellow")
+    table.add_column("Last Checked")
+    table.add_column("Max Token")
+    table.add_column("Concurrency")
+    table.add_column("Sources", justify="right")
+
+    for key, info in data["keys"].items():
+        if args.provider and info["provider"].lower() != args.provider.lower():
+            continue
+        if args.status and info["status"] != args.status:
+            continue
+        if args.batch:
+            has_batch = any(s.get("batch") == args.batch for s in info.get("sources", []))
+            if not has_batch:
+                continue
+
+        status_style = {
+            "valid": "[green]valid[/green]",
+            "invalid": "[red]invalid[/red]",
+            "error": "[red]error[/red]",
+        }.get(info["status"], f"[yellow]{info['status']}[/yellow]")
+
+        tests = info.get("tests", {})
+        max_tokens = str(tests.get("max_tokens", "-")) if tests.get("max_tokens") else "-"
+        max_conc = str(tests.get("max_concurrency", "-")) if tests.get("max_concurrency") else "-"
+
+        table.add_row(
+            info["key_masked"],
+            info["provider"],
+            status_style,
+            info.get("last_checked", "never") or "never",
+            max_tokens,
+            max_conc,
+            str(len(info.get("sources", [])))
+        )
+
+    console.print(table)
+
+
+def cmd_report(args, config):
+    keys_path = Path(config["storage"]["keys_file"])
+    if not keys_path.exists():
+        console.print("[yellow]No keys file found.[/yellow]")
+        return
+
+    try:
+        data = _load_keys(config)
+    except KeyManagerError as e:
+        console.print(f"[red]Error:[/red] {e.message}")
+        return
+
+    days = args.days or 7
+    cutoff = datetime.utcnow() - timedelta(days=days)
+
+    stats = {}
+    for key, info in data["keys"].items():
+        provider = info["provider"]
+        if provider not in stats:
+            stats[provider] = {"total": 0, "valid": 0, "invalid": 0, "error": 0, "max_tokens": [], "concurrency": []}
+        stats[provider]["total"] += 1
+        status = info.get("status", "unknown")
+        if status in stats[provider]:
+            stats[provider][status] += 1
+        tests = info.get("tests", {})
+        if tests.get("max_tokens"):
+            stats[provider]["max_tokens"].append(tests["max_tokens"])
+        if tests.get("max_concurrency"):
+            stats[provider]["concurrency"].append(tests["max_concurrency"])
+
+    table = Table(title=f"Key Statistics (last {days} days)")
+    table.add_column("Provider", style="green")
+    table.add_column("Total", justify="right")
+    table.add_column("Valid", justify="right", style="green")
+    table.add_column("Invalid", justify="right", style="red")
+    table.add_column("Error", justify="right", style="red")
+    table.add_column("MaxToken")
+    table.add_column("Concurrency")
+
+    total_all = 0
+    total_valid = 0
+    total_invalid = 0
+    total_error = 0
+
+    for provider, s in sorted(stats.items()):
+        total_all += s["total"]
+        total_valid += s["valid"]
+        total_invalid += s["invalid"]
+        total_error += s["error"]
+
+        max_tokens = max(s["max_tokens"]) if s["max_tokens"] else "-"
+        max_conc = max(s["concurrency"]) if s["concurrency"] else "-"
+
+        table.add_row(
+            provider,
+            str(s["total"]),
+            str(s["valid"]),
+            str(s["invalid"]),
+            str(s["error"]),
+            str(max_tokens),
+            str(max_conc)
+        )
+
+    table.add_section()
+    table.add_row(
+        "TOTAL",
+        str(total_all),
+        str(total_valid),
+        str(total_invalid),
+        str(total_error),
+        "-", "-"
+    )
+
+    console.print(table)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="API Key Manager")
+    subparsers = parser.add_subparsers(dest="command")
+
+    # Import command
+    import_parser = subparsers.add_parser("import", help="Import API keys")
+    import_parser.add_argument("--file", help="JSON file to import")
+    import_parser.add_argument("--dir", help="Directory to scan")
+    import_parser.add_argument("--batch", help="Batch label")
+
+    # Check command
+    check_parser = subparsers.add_parser("check", help="Validate keys")
+    check_parser.add_argument("--provider", help="Filter by provider")
+    check_parser.add_argument("--status", help="Filter by status")
+    check_parser.add_argument("--key", help="Check single key")
+
+    # Test command
+    test_parser = subparsers.add_parser("test", help="Token/concurrency test")
+    test_parser.add_argument("--skip-token", action="store_true", help="Skip token test")
+    test_parser.add_argument("--skip-concurrency", action="store_true", help="Skip concurrency test")
+    test_parser.add_argument("--provider", help="Filter by provider")
+    test_parser.add_argument("--key", help="Test single key")
+
+    # List command
+    list_parser = subparsers.add_parser("list", help="List keys")
+    list_parser.add_argument("--provider", help="Filter by provider")
+    list_parser.add_argument("--status", help="Filter by status")
+    list_parser.add_argument("--batch", help="Filter by batch")
+
+    # Report command
+    report_parser = subparsers.add_parser("report", help="Show statistics")
+    report_parser.add_argument("--days", type=int, help="Days to include")
+
+    args = parser.parse_args()
+    config = load_config()
+
+    if args.command == "import":
+        cmd_import(args, config)
+    elif args.command == "check":
+        cmd_check(args, config)
+    elif args.command == "test":
+        cmd_test(args, config)
+    elif args.command == "list":
+        cmd_list(args, config)
+    elif args.command == "report":
+        cmd_report(args, config)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

key_manager/config.py

@@ -0,0 +1,61 @@
+import yaml
+from pathlib import Path
+
+DEFAULT_CONFIG = {
+    "scan": {
+        "directories": ["./data/input"],
+        "recursive": True,
+    },
+    "check": {
+        "interval_hours": 6,
+        "timeout_seconds": 30,
+        "concurrency": 100,
+        "retry_failed": True,
+        "retry_count": 2,
+    },
+    "test": {
+        "token_test": True,
+        "token_auto_detect": True,
+        "token_steps": [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576],
+        "token_max_manual": None,
+        "concurrency_test": True,
+        "concurrency_steps": [1, 5, 10, 20, 50, 100],
+        "concurrency_timeout_seconds": 120,
+    },
+    "storage": {
+        "keys_file": "./data/keys.json",
+        "check_results_file": "./data/check_results.json",
+        "test_results_file": "./data/test_results.json",
+        "logs_dir": "./data/logs",
+    },
+    "auth": {
+        "api_key": "",
+    },
+    "rate_limit": {
+        "requests_per_minute": 60,
+    },
+}
+
+
+def _deep_merge(base: dict, override: dict) -> dict:
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = _deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def load_config(path: str = "config.yaml") -> dict:
+    config_path = Path(path)
+    if config_path.exists():
+        with open(config_path, "r", encoding="utf-8") as f:
+            user_config = yaml.safe_load(f) or {}
+        return _deep_merge(DEFAULT_CONFIG, user_config)
+    return DEFAULT_CONFIG.copy()
+
+
+def save_config(config: dict, path: str = "config.yaml"):
+    with open(path, "w", encoding="utf-8") as f:
+        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)

key_manager/core.py

@@ -0,0 +1,205 @@
+"""High-level facade for programmatic usage."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional
+
+from key_manager.config import load_config
+from key_manager.storage import KeyStore
+from key_manager.providers import PROVIDERS, get_display_name
+from key_manager.detector import detect_provider
+
+
+class KeyManager:
+    """Convenient facade for managing API keys programmatically.
+    
+    Usage:
+        km = KeyManager()
+        km.import_keys("keys.json")
+        results = km.check(provider="openai")
+        km.save()
+    """
+    
+    def __init__(
+        self, 
+        config_path: str = "config.yaml", 
+        config: dict[str, Any] | None = None
+    ):
+        """Initialize KeyManager.
+        
+        Args:
+            config_path: Path to config.yaml file
+            config: Direct config dict (overrides config_path)
+        """
+        self.config = config or load_config(config_path)
+        self._store = KeyStore(
+            self.config["storage"]["keys_file"], 
+            self.config
+        )
+    
+    @property
+    def providers(self) -> dict[str, Any]:
+        """Available provider registry."""
+        return PROVIDERS
+    
+    def list_providers(self) -> list[str]:
+        """List available provider names."""
+        return list(PROVIDERS.keys())
+    
+    def get_provider(self, name: str):
+        """Get a provider by name."""
+        return PROVIDERS.get(name)
+    
+    def get_provider_display_name(self, name: str) -> str:
+        """Get human-readable display name for a provider."""
+        return get_display_name(name)
+    
+    def load_keys(self) -> dict[str, Any]:
+        """Load keys from encrypted storage."""
+        return self._store.load()
+    
+    def save_keys(self, data: dict[str, Any]) -> None:
+        """Save keys to encrypted storage."""
+        self._store.save(data)
+    
+    def detect_provider(self, key: str) -> list[str]:
+        """Auto-detect provider from key prefix."""
+        return detect_provider(key)
+    
+    def list_keys(
+        self, 
+        provider: str | None = None,
+        status: str | None = None
+    ) -> dict[str, Any]:
+        """List keys with optional filtering.
+        
+        Args:
+            provider: Filter by provider name
+            status: Filter by status (valid/invalid/error)
+            
+        Returns:
+            Dict with keys matching filters
+        """
+        data = self.load_keys()
+        result = {}
+        
+        for key, info in data.get("keys", {}).items():
+            if provider and info.get("provider", "").lower() != provider.lower():
+                continue
+            if status and info.get("status") != status:
+                continue
+            result[key] = info
+        
+        return result
+    
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about stored keys.
+        
+        Returns:
+            Dict with provider counts and status breakdown
+        """
+        data = self.load_keys()
+        stats: dict[str, Any] = {
+            "total": 0,
+            "by_provider": {},
+            "by_status": {"valid": 0, "invalid": 0, "error": 0, "unknown": 0}
+        }
+        
+        for key, info in data.get("keys", {}).items():
+            stats["total"] += 1
+            provider = info.get("provider", "unknown")
+            status = info.get("status", "unknown")
+            
+            if provider not in stats["by_provider"]:
+                stats["by_provider"][provider] = 0
+            stats["by_provider"][provider] += 1
+            
+            if status in stats["by_status"]:
+                stats["by_status"][status] += 1
+        
+        return stats
+    
+    async def check_key(
+        self, 
+        key: str, 
+        provider: str | None = None,
+        timeout: int = 30
+    ) -> dict[str, Any]:
+        """Check a single key.
+        
+        Args:
+            key: API key to check
+            provider: Provider name (auto-detected if not specified)
+            timeout: Request timeout in seconds
+            
+        Returns:
+            Dict with check results
+        """
+        import httpx
+        from key_manager.providers import PROVIDERS
+        
+        if not provider:
+            providers = self.detect_provider(key)
+            if not providers:
+                return {"valid": False, "error": "Could not detect provider"}
+            provider = providers[0]
+        
+        provider_obj = PROVIDERS.get(provider)
+        if not provider_obj:
+            return {"valid": False, "error": f"Unknown provider: {provider}"}
+        
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            result = await provider_obj.check(client, key)
+            return {
+                "valid": result.valid,
+                "provider": provider,
+                "status_code": result.status_code,
+                "latency_ms": result.latency_ms,
+                "error": result.error
+            }
+    
+    async def check_all(
+        self,
+        provider: str | None = None,
+        status: str | None = None,
+        concurrency: int = 100,
+        timeout: int = 30
+    ) -> dict[str, Any]:
+        """Check all keys in storage.
+        
+        Args:
+            provider: Filter by provider
+            status: Filter by status
+            concurrency: Max concurrent checks
+            timeout: Request timeout
+            
+        Returns:
+            Dict with check results
+        """
+        from key_manager.checker import run_check
+        
+        return await run_check(
+            keys_file=self.config["storage"]["keys_file"],
+            results_file=self.config["storage"]["check_results_file"],
+            logs_dir=self.config["storage"]["logs_dir"],
+            concurrency=concurrency,
+            timeout=timeout,
+            proxy=self.config.get("proxy") or None,
+            retry_failed=self.config["check"]["retry_failed"],
+            retry_count=self.config["check"]["retry_count"]
+        )
+    
+    async def validate_keys(
+        self,
+        provider: str | None = None,
+        status: str | None = None,
+        concurrency: int = 100,
+        timeout: int = 30
+    ) -> dict[str, Any]:
+        """Validate all keys (alias for check_all)."""
+        return await self.check_all(
+            provider=provider,
+            status=status,
+            concurrency=concurrency,
+            timeout=timeout
+        )