new_api_proxy.py

import os
from urllib.parse import urlsplit, urlunsplit


NEW_API_PROXY_URL_ENV = "NEW_API_PROXY_URL"
NEW_API_PROXY_TOKEN_ENV = "NEW_API_PROXY_TOKEN"
NEW_API_PROXY_AUTH_QUERY_ENV = "NEW_API_PROXY_AUTH_QUERY"
NEW_API_PROXY_DIRECT_FALLBACK_ENV = "NEW_API_PROXY_DIRECT_FALLBACK"
NEW_API_PROXY_TOKEN_HEADER = "X-New-API-Proxy-Token"


def get_new_api_proxy_base_url(proxy_url=None):
    raw = proxy_url if proxy_url is not None else os.getenv(NEW_API_PROXY_URL_ENV, "")
    raw = str(raw or "").strip().strip('"').strip("'")
    if not raw:
        return ""
    if "://" not in raw:
        raw = f"https://{raw}"
    return raw.rstrip("/")


def build_new_api_proxy_url(path, proxy_url=None):
    proxy_base_url = get_new_api_proxy_base_url(proxy_url)
    if not proxy_base_url:
        return ""

    path = str(path or "")
    if not path.startswith("/"):
        path = f"/{path}"

    base_parts = urlsplit(proxy_base_url)
    path_part, _, path_query = path.partition("?")
    base_path = base_parts.path.rstrip("/")
    if base_path and base_path != "/" and not base_path.endswith(path_part.rstrip("/")):
        final_path = f"{base_path}{path_part}"
    elif base_path and base_path.endswith(path_part.rstrip("/")):
        final_path = base_path
    else:
        final_path = path_part

    extra_query = os.getenv(NEW_API_PROXY_AUTH_QUERY_ENV, "").strip().lstrip("?")
    query_parts = [q for q in (base_parts.query, extra_query, path_query) if q]
    return urlunsplit(
        (
            base_parts.scheme,
            base_parts.netloc,
            final_path,
            "&".join(query_parts),
            base_parts.fragment,
        )
    )


def with_new_api_proxy_headers(headers=None, proxy_url=None):
    headers = dict(headers or {})
    if get_new_api_proxy_base_url(proxy_url):
        proxy_token = os.getenv(NEW_API_PROXY_TOKEN_ENV, "").strip()
        if proxy_token:
            headers[NEW_API_PROXY_TOKEN_HEADER] = proxy_token
    return headers


def new_api_proxy_direct_fallback_enabled():
    raw = os.getenv(NEW_API_PROXY_DIRECT_FALLBACK_ENV, "1").strip().lower()
    return raw not in {"0", "false", "no", "off"}