key_manager.py

#!/usr/bin/env python3
"""
FreeBG API Key Manager
Use this script locally to manage customer API keys.
Run: python key_manager.py
"""

import json
import random
import string
import os
from datetime import datetime

# ── Plans ──────────────────────────────────────────────────────────────────────
PLANS = {
    "free":    {"daily": 10,    "models": ["fast"],                   "price": "$0"},
    "starter": {"daily": 100,   "models": ["fast", "quality"],        "price": "$9/mo"},
    "pro":     {"daily": 500,   "models": ["fast", "quality", "best"],"price": "$29/mo"},
    "master":  {"daily": 999999,"models": ["fast", "quality", "best"],"price": "Custom"},
}

KEYS_FILE = "api_keys.json"


def load_keys():
    if os.path.exists(KEYS_FILE):
        with open(KEYS_FILE) as f:
            return json.load(f)
    return {}


def save_keys(keys):
    with open(KEYS_FILE, "w") as f:
        json.dump(keys, f, indent=2)
    print(f"✅ Saved to {KEYS_FILE}")


def generate_key(plan: str) -> str:
    chars = string.ascii_lowercase + string.digits
    part1 = ''.join(random.choices(chars, k=8))
    part2 = ''.join(random.choices(chars, k=8))
    return f"freebg-{plan[:2]}-{part1}-{part2}"


def add_customer(owner: str, plan: str, custom_limit: int = None):
    keys = load_keys()
    
    key = generate_key(plan)
    plan_config = PLANS.get(plan, PLANS["free"])
    
    keys[key] = {
        "plan": plan,
        "owner": owner,
        "calls_today": 0,
        "reset_at": 0,
        "limit": custom_limit or plan_config["daily"],
        "created_at": datetime.utcnow().isoformat(),
        "models": plan_config["models"],
    }
    
    save_keys(keys)
    
    print(f"""
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
✅ New API Key Created
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Owner:  {owner}
Plan:   {plan} ({plan_config['price']})
Key:    {key}
Limit:  {custom_limit or plan_config['daily']} calls/day
Models: {', '.join(plan_config['models'])}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Send this key to your customer:
  {key}

Now update your HuggingFace Space:
  Settings → Secrets → API_KEYS_JSON
  Paste the content of {KEYS_FILE}
""")
    return key


def list_customers():
    keys = load_keys()
    if not keys:
        print("No customers yet.")
        return
    
    print(f"\n{'Key':<40} {'Owner':<25} {'Plan':<10} {'Calls'}")
    print("─" * 90)
    for key, data in keys.items():
        print(f"{key:<40} {data.get('owner',''):<25} {data.get('plan',''):<10} {data.get('calls_today', 0)}")


def revoke_key(key: str):
    keys = load_keys()
    if key in keys:
        owner = keys[key].get("owner", "")
        del keys[key]
        save_keys(keys)
        print(f"🗑️  Revoked key for: {owner}")
    else:
        print("Key not found.")


def export_for_hf():
    """Export API_KEYS_JSON for HuggingFace Space secret"""
    keys = load_keys()
    print("\n📋 Copy this value to HuggingFace Space Secret 'API_KEYS_JSON':\n")
    print(json.dumps(keys))
    print()


def upgrade_plan(key: str, new_plan: str):
    keys = load_keys()
    if key not in keys:
        print("Key not found.")
        return
    
    plan_config = PLANS.get(new_plan, PLANS["free"])
    keys[key]["plan"] = new_plan
    keys[key]["limit"] = plan_config["daily"]
    keys[key]["models"] = plan_config["models"]
    save_keys(keys)
    
    owner = keys[key].get("owner", "")
    print(f"⬆️  Upgraded {owner} to {new_plan} plan ({plan_config['daily']} calls/day)")


# ── CLI ────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys
    
    print("""
╔══════════════════════════════════════╗
║   FreeBG API Key Manager             ║
╚══════════════════════════════════════╝
    """)
    
    if len(sys.argv) < 2:
        print("Commands:")
        print("  python key_manager.py add customer@gmail.com starter     - Add new customer")
        print("  python key_manager.py add <email> <plan>     - Add new customer")
        print("  python key_manager.py list                   - List all customers")
        print("  python key_manager.py revoke <key>           - Revoke a key")
        print("  python key_manager.py upgrade <key> <plan>   - Upgrade plan")
        print("  python key_manager.py export                 - Export for HF Space")
        print()
        print("Plans:", ", ".join(PLANS.keys()))
        sys.exit(0)
    
    cmd = sys.argv[1]
    
    if cmd == "add":
        if len(sys.argv) < 4:
            print("Usage: python key_manager.py add <email> <plan>")
        else:
            add_customer(sys.argv[2], sys.argv[3])
    
    elif cmd == "list":
        list_customers()
    
    elif cmd == "revoke":
        if len(sys.argv) < 3:
            print("Usage: python key_manager.py revoke <key>")
        else:
            revoke_key(sys.argv[2])
    
    elif cmd == "upgrade":
        if len(sys.argv) < 4:
            print("Usage: python key_manager.py upgrade <key> <plan>")
        else:
            upgrade_plan(sys.argv[2], sys.argv[3])
    
    elif cmd == "export":
        export_for_hf()
    
    else:
        print(f"Unknown command: {cmd}")