key_manager.py

#!/usr/bin/env python3
"""
FreeME Magic Eraser — Key Manager
Same pattern as FreeBG background remover.

Usage:
  python key_manager.py add customer@email.com starter
  python key_manager.py list
  python key_manager.py upgrade freebg-me-st-abc12345-xyz67890 pro
  python key_manager.py revoke freebg-me-st-abc12345-xyz67890
  python key_manager.py export
"""

import json
import sys
import secrets
import datetime
import os

DB_FILE = "keys_db.json"

PLANS = {
    "free":    {"daily": 10,   "models": ["fast"]},
    "starter": {"daily": 100,  "models": ["fast", "quality"]},
    "pro":     {"daily": 500,  "models": ["fast", "quality", "best"]},
    "master":  {"daily": 99999,"models": ["fast", "quality", "best"]},
}

PREFIX = {
    "free":    "me-fr",
    "starter": "me-st",
    "pro":     "me-pr",
    "master":  "me-ma",
}

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

def save_db(db):
    with open(DB_FILE, "w") as f:
        json.dump(db, f, indent=2)

def today():
    return datetime.datetime.utcnow().strftime("%Y-%m-%d")

def cmd_add(email, plan):
    if plan not in PLANS:
        print(f"❌ Invalid plan. Choose: {', '.join(PLANS.keys())}")
        return
    db = load_db()
    pre = PREFIX[plan]
    key = f"freebg-{pre}-{secrets.token_hex(4)}-{secrets.token_hex(4)}"
    db[key] = {
        "owner": email,
        "plan": plan,
        "created": today(),
        "active": True,
    }
    save_db(db)
    print(f"\n✅ Key created:")
    print(f"   Key   : {key}")
    print(f"   Owner : {email}")
    print(f"   Plan  : {plan} ({PLANS[plan]['daily']}/day, models: {PLANS[plan]['models']})")
    print(f"\n📋 Next step: run 'python key_manager.py export' then update HF Space Secret\n")

def cmd_list():
    db = load_db()
    if not db:
        print("No keys found.")
        return
    print(f"\n{'Key':<46} {'Owner':<28} {'Plan':<10} {'Created':<12} {'Active'}")
    print("-" * 105)
    for k, v in db.items():
        print(f"{k:<46} {v.get('owner',''):<28} {v.get('plan','?'):<10} {v.get('created',''):<12} {v.get('active', True)}")
    print()

def cmd_upgrade(key, new_plan):
    if new_plan not in PLANS:
        print(f"❌ Invalid plan. Choose: {', '.join(PLANS.keys())}")
        return
    db = load_db()
    if key not in db:
        print(f"❌ Key not found: {key}")
        return
    old_plan = db[key].get("plan", "?")
    db[key]["plan"] = new_plan
    save_db(db)
    print(f"✅ Upgraded {key}")
    print(f"   {old_plan} → {new_plan}")
    print(f"\n📋 Run 'python key_manager.py export' then update HF Space Secret\n")

def cmd_revoke(key):
    db = load_db()
    if key not in db:
        print(f"❌ Key not found: {key}")
        return
    owner = db[key].get("owner", "")
    del db[key]
    save_db(db)
    print(f"✅ Key revoked: {key} (was: {owner})")
    print(f"\n📋 Run 'python key_manager.py export' then update HF Space Secret\n")

def cmd_export():
    db = load_db()
    exported = json.dumps(db, indent=2)
    print("\n" + "="*60)
    print("COPY THE JSON BELOW → HF Space → Settings → Secrets → API_KEYS_JSON → Replace")
    print("="*60)
    print(exported)
    print("="*60 + "\n")

def main():
    args = sys.argv[1:]
    if not args:
        print(__doc__)
        return

    cmd = args[0].lower()

    if cmd == "add" and len(args) == 3:
        cmd_add(args[1], args[2])
    elif cmd == "list":
        cmd_list()
    elif cmd == "upgrade" and len(args) == 3:
        cmd_upgrade(args[1], args[2])
    elif cmd == "revoke" and len(args) == 2:
        cmd_revoke(args[1])
    elif cmd == "export":
        cmd_export()
    else:
        print(__doc__)

if __name__ == "__main__":
    main()