Custom gateway provisioners

#!/usr/bin/env python3
import hashlib
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

def respond(payload, exit_code=0):
    json.dump(payload, sys.stdout, separators=(",", ":"))
    sys.stdout.write("\n")
    raise SystemExit(exit_code)

class ProvisionerError(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message

def required_env(name):
    value = os.environ.get(name, "").strip()
    if not value:
        raise ProvisionerError("invalid_config", f"{name} is required")
    return value

BASE_URL = required_env("HARNESS_GATEWAY_URL").rstrip("/")
ADMIN_KEY = required_env("HARNESS_LITELLM_ADMIN_KEY")
MODELS = [value for value in os.environ.get("LITELLM_MODELS", "").split(",") if value]

def litellm(method, path, body=None, query=None):
    url = f"{BASE_URL}/{path.lstrip('/')}"
    if query:
        url += "?" + urllib.parse.urlencode(query)
    encoded = None if body is None else json.dumps(body).encode()
    request = urllib.request.Request(url, data=encoded, method=method, headers={
        "Authorization": f"Bearer {ADMIN_KEY}",
        "Accept": "application/json",
        "Content-Type": "application/json",
    })
    try:
        with urllib.request.urlopen(request, timeout=20) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        # Do not include the upstream body: it can contain credentials.
        if path == "/key/info" and error.code == 404:
            raise ProvisionerError("credential_invalid", "LiteLLM key was deleted")
        raise ProvisionerError("rejected", f"LiteLLM returned HTTP {error.code}")
    except (urllib.error.URLError, TimeoutError):
        raise ProvisionerError("unavailable", "LiteLLM request failed")
    except (UnicodeDecodeError, json.JSONDecodeError):
        raise ProvisionerError("rejected", "LiteLLM returned invalid JSON")

def find_user(email):
    payload = litellm("GET", "/user/list", query={
        "user_email": email,
        "page_size": 100,
    })
    users = [user for user in payload.get("users", [])
             if user.get("user_email", "").casefold() == email.casefold()]
    if not users:
        raise ProvisionerError("account_missing", email)
    if len(users) != 1:
        raise ProvisionerError("conflict", f"multiple LiteLLM users match {email}")
    user = users[0]
    if not user.get("user_id"):
        raise ProvisionerError("rejected", "LiteLLM user has no user_id")
    teams = [team.get("team_id", team.get("id"))
             for team in user.get("teams", [])]
    return user["user_id"], [team for team in teams if team]

def ensure(request):
    identity = request["identity"]
    user_id, teams = find_user(identity["email"])
    alias = f"blue:{identity['email'].strip().lower()}"
    previous = request.get("previous")
    metadata = {
        "managed_by": "blue",
        "owner_email": identity["email"].strip().lower(),
        "provisioner": "organization-litellm",
    }
    key_policy = {
        "user_id": user_id,
        "key_alias": alias,
        "models": MODELS,
        "metadata": metadata,
    }

    # Reconcile an existing virtual key without returning plaintext again.
    if previous:
        current = litellm("GET", "/key/info", query={"key": previous["external_id"]})
        if current.get("info", {}).get("blocked") is True:
            raise ProvisionerError("credential_invalid", "LiteLLM key was blocked")
        litellm("POST", "/key/update", {
            **key_policy,
            "key": previous["external_id"],
        })
        return {
            "credential": None,
            "external_id": previous["external_id"],
            "alias": alias,
            "metadata": {"team_ids": teams, "models": MODELS},
            "expires_at": None,
        }

    generated = litellm("POST", "/key/generate", {
        **key_policy,
        "key_type": "llm_api",
    })
    credential = generated.get("key")
    if not credential:
        raise ProvisionerError("rejected", "LiteLLM generated no key")
    # LiteLLM commonly returns the stored token/hash separately. Fall back to
    # a SHA-256 identifier so plaintext is never persisted as external_id.
    external_id = generated.get("token") or hashlib.sha256(
        credential.encode()
    ).hexdigest()
    return {
        "credential": credential,
        "external_id": external_id,
        "alias": alias,
        "metadata": {"team_ids": teams, "models": MODELS},
        "expires_at": generated.get("expires"),
    }

def revoke(request):
    litellm("POST", "/key/delete", {"keys": [request["external_id"]]})
    return {"revoked": True}

try:
    envelope = json.load(sys.stdin)
    if envelope.get("protocol_version") != 1:
        raise ProvisionerError("invalid_config", "unsupported protocol version")
    operation = envelope.get("operation")
    if operation == "ensure":
        result = ensure(envelope["request"])
    elif operation == "revoke":
        result = revoke(envelope["request"])
    else:
        raise ProvisionerError("invalid_config", "unsupported operation")
    respond({"protocol_version": 1, "status": "success", "result": result})
except ProvisionerError as error:
    respond({"protocol_version": 1, "status": "error", "error": {
        "code": error.code, "message": error.message
    }}, 1)
except (KeyError, TypeError, json.JSONDecodeError):
    respond({"protocol_version": 1, "status": "error", "error": {
        "code": "invalid_config", "message": "invalid provisioner request"
    }}, 1)