autoreg-skills-max/skills/account-validation/templates/openai_key_pool_proxy.py

#!/usr/bin/env python3
"""OpenAI-compatible key-pool proxy template.

Use this only with legitimate upstream API keys supplied by the user, e.g.
OPENAI_KEYS=sk-proj-xxx,sk-proj-yyy. Do not adapt it to ChatGPT web
session tokens, Auth0 access tokens, cookies, or bulk signup/token pooling.

Run:
  OPENAI_KEYS=sk-proj-xxx PROXY_AUTH_TOKEN=local-token python openai_key_pool_proxy.py

Test:
  curl http://127.0.0.1:8091/health
  curl http://127.0.0.1:8091/v1/models -H 'Authorization: Bearer local-token'
"""

from __future__ import annotations

import json
import os
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

UPSTREAM_BASE = os.getenv("UPSTREAM_BASE", "https://api.openai.com").rstrip("/")
PROXY_AUTH_TOKEN = os.getenv("PROXY_AUTH_TOKEN", "local-token")


def redact(key: str) -> str:
    return key[:7] + "..." + key[-4:] if len(key) > 12 else key[:4] + "..."


@dataclass
class KeyState:
    key: str
    name: str
    enabled: bool = True
    last_status: int | None = None
    last_error: str | None = None
    ok_count: int = 0
    fail_count: int = 0
    last_used: float | None = None
    disabled_reason: str | None = None

    def public(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "key": redact(self.key),
            "enabled": self.enabled,
            "last_status": self.last_status,
            "last_error": self.last_error,
            "ok_count": self.ok_count,
            "fail_count": self.fail_count,
            "last_used": self.last_used,
            "disabled_reason": self.disabled_reason,
        }


@dataclass
class Pool:
    keys: list[KeyState]
    cursor: int = 0
    lock: threading.Lock = field(default_factory=threading.Lock)

    def choose(self) -> KeyState | None:
        with self.lock:
            live = [k for k in self.keys if k.enabled]
            if not live:
                return None
            item = live[self.cursor % len(live)]
            self.cursor += 1
            item.last_used = time.time()
            return item

    def mark(self, item: KeyState, status: int, body: bytes) -> None:
        text = body[:1200].decode("utf-8", errors="ignore")
        with self.lock:
            item.last_status = status
            item.last_error = None if 200 <= status < 300 else text
            if 200 <= status < 300:
                item.ok_count += 1
                return
            item.fail_count += 1
            if status in (401, 403):
                item.enabled = False
                item.disabled_reason = f"HTTP {status}"
            elif status == 429 and "billing_not_active" in text:
                item.enabled = False
                item.disabled_reason = "billing_not_active"

    def public(self) -> dict[str, Any]:
        with self.lock:
            return {"total": len(self.keys), "enabled": sum(k.enabled for k in self.keys), "keys": [k.public() for k in self.keys]}


def load_keys() -> Pool:
    raw = os.getenv("OPENAI_KEYS", "").replace("\n", ",")
    seen: set[str] = set()
    items: list[KeyState] = []
    for idx, key in enumerate([x.strip() for x in raw.split(",") if x.strip()], 1):
        if key in seen:
            continue
        seen.add(key)
        items.append(KeyState(key=key, name=f"key{idx}"))
    return Pool(items)


POOL = load_keys()


class Handler(BaseHTTPRequestHandler):
    def _json(self, status: int, obj: Any) -> None:
        body = json.dumps(obj, ensure_ascii=False, indent=2).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _authorized(self) -> bool:
        return not PROXY_AUTH_TOKEN or self.headers.get("Authorization") == f"Bearer {PROXY_AUTH_TOKEN}"

    def _body(self) -> bytes:
        return self.rfile.read(int(self.headers.get("Content-Length") or 0))

    def do_GET(self) -> None:  # noqa: N802
        if self.path in ("/", "/health"):
            return self._json(200, {"ok": True, "upstream": UPSTREAM_BASE, "pool": POOL.public()})
        return self._proxy()

    def do_POST(self) -> None:  # noqa: N802
        return self._proxy()

    def _proxy(self) -> None:
        if not self.path.startswith("/v1/"):
            return self._json(404, {"error": "only /v1/* is proxied"})
        if not self._authorized():
            return self._json(401, {"error": "proxy auth required"})
        body = self._body()
        attempts = max(1, min(3, len([k for k in POOL.keys if k.enabled]) or 1))
        last_status, last_body, last_headers = 503, b'{"error":"no usable key"}', {"Content-Type": "application/json"}
        for _ in range(attempts):
            item = POOL.choose()
            if not item:
                break
            status, resp_body, resp_headers = call_upstream(self.command, self.path, self.headers, body, item.key)
            POOL.mark(item, status, resp_body)
            last_status, last_body, last_headers = status, resp_body, resp_headers
            if status < 400 or status in (400, 404, 422):
                break
            if status not in (401, 403, 429, 500, 502, 503, 504):
                break
        self.send_response(last_status)
        self.send_header("Content-Type", last_headers.get("Content-Type", "application/json"))
        self.send_header("Content-Length", str(len(last_body)))
        self.end_headers()
        self.wfile.write(last_body)


def call_upstream(method: str, path: str, incoming_headers: Any, body: bytes, upstream_key: str) -> tuple[int, bytes, dict[str, str]]:
    headers = {k: v for k, v in incoming_headers.items() if k.lower() not in {"host", "authorization", "content-length", "connection"}}
    headers["Authorization"] = f"Bearer {upstream_key}"
    if body and "Content-Type" not in headers:
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(UPSTREAM_BASE + path, data=body if method != "GET" else None, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            return resp.status, resp.read(), dict(resp.headers.items())
    except urllib.error.HTTPError as e:
        return e.code, e.read(), dict(e.headers.items())
    except Exception as e:
        return 502, json.dumps({"error": str(e)}, ensure_ascii=False).encode(), {"Content-Type": "application/json"}


if __name__ == "__main__":
    print(f"proxy=http://127.0.0.1:8091 upstream={UPSTREAM_BASE} keys={len(POOL.keys)}")
    ThreadingHTTPServer(("127.0.0.1", 8091), Handler).serve_forever()