autoreg-skills-max: initial public release (restored, redacted, hardened)

skills/account-validation/templates/openai_key_pool_proxy.py

@@ -0,0 +1,180 @@
+#!/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()