truf-server/app/keycheckers/provider_resolution.py

import sys

sys.dont_write_bytecode = True

import os


SUPPORTED_PROVIDERS = ("deepseek", "zai", "qwen", "kimi")
DEFAULT_PROVIDER_ORDER = SUPPORTED_PROVIDERS
AMBIGUOUS_QWEN_DEEPSEEK_HINT = "ambiguous_qwen_deepseek"
AMBIGUOUS_GENERIC_SK_HINT = "ambiguous_generic_sk"


def split_csv(value):
    if not value:
        return []
    if isinstance(value, str):
        values = value.split(",")
    else:
        values = value
    return [str(item).strip().lower() for item in values if str(item).strip()]


def unique_supported(values):
    output = []
    seen = set()
    for value in values:
        provider = str(value or "").strip().lower()
        if provider in SUPPORTED_PROVIDERS and provider not in seen:
            seen.add(provider)
            output.append(provider)
    return output


def providers_for_hint(hint):
    hint = str(hint or "").strip().lower()
    if hint == AMBIGUOUS_QWEN_DEEPSEEK_HINT:
        return ["qwen", "deepseek"]
    if hint == AMBIGUOUS_GENERIC_SK_HINT:
        return list(SUPPORTED_PROVIDERS)
    return [hint] if hint in SUPPORTED_PROVIDERS else []


def detector_provider(finding):
    if not isinstance(finding, dict):
        return ""
    extra = finding.get("ExtraData") if isinstance(finding.get("ExtraData"), dict) else {}
    names = {
        str(finding.get("DetectorName") or finding.get("DetectorType") or "").strip().lower(),
        str(extra.get("name") or "").strip().lower(),
    }
    mappings = (
        ("deepseek", {"deepseek", "deepseekapikey", "deepseek_api_key"}),
        ("zai", {"zaiglm"}),
        ("qwen", {"qwendashscope", "qwen_dashscope", "qwen", "dashscope"}),
        ("kimi", {"kimimoonshot", "moonshotai", "moonshot", "kimi"}),
    )
    for provider, detectors in mappings:
        if names & detectors:
            return provider
    return ""


def ordered_providers(finding=None, hint="", origin_service="", configured_order=None):
    context = finding.get("ScannerContext") if isinstance(finding, dict) and isinstance(
        finding.get("ScannerContext"), dict
    ) else {}
    hint = str(hint or context.get("provider_hint") or "").strip().lower()
    compatible = unique_supported(context.get("provider_candidates") or providers_for_hint(hint))
    if not compatible:
        compatible = providers_for_hint(hint)
    if not compatible:
        compatible = list(SUPPORTED_PROVIDERS)

    configured = unique_supported(
        configured_order
        if configured_order is not None
        else split_csv(os.getenv("KEYCHECK_PROVIDER_RESOLUTION_ORDER"))
    )
    base_order = configured or list(DEFAULT_PROVIDER_ORDER)
    origin = str(origin_service or detector_provider(finding)).strip().lower()
    ordered = []
    if origin in compatible:
        ordered.append(origin)
    ordered.extend(provider for provider in base_order if provider in compatible)
    ordered.extend(provider for provider in compatible if provider not in ordered)
    return unique_supported(ordered)


def provider_result_outcome(result):
    result = result if isinstance(result, dict) else {}
    status = str(result.get("status") or "UNKNOWN").strip().upper()
    if result.get("authenticated") is True or status in ("VALID", "ALIVE"):
        return "match"
    if result.get("candidate_rejected") or status in (
        "DEAD", "INVALID", "EXPIRED", "LEAKED_REVOKED", "INVALID_OR_REVOKED",
    ):
        return "no_match"
    return "retry"


def bounded_attempt(provider, result, outcome):
    result = result if isinstance(result, dict) else {}
    error = result.get("error") if isinstance(result.get("error"), dict) else {}
    return {
        "provider": provider,
        "outcome": outcome,
        "status": str(result.get("status") or "UNKNOWN").upper(),
        "authenticated": bool(result.get("authenticated")),
        "http_status": int(result.get("http_status") or error.get("http_status") or 0),
        "business_code": str(result.get("business_code") or error.get("code") or "")[:80],
        "region": str(result.get("region") or "")[:160],
        "message": str(result.get("message") or "").replace("\r", " ").replace("\n", " ")[:300],
    }


def default_provider_probe(provider, key, proxy, timeout, debug=False):
    if provider == "deepseek":
        from keycheckers.deepseek import deepseekKeycheck

        return deepseekKeycheck.check_key(key, proxy, timeout)
    if provider == "zai":
        from keycheckers.zai import zaiKeycheck

        return zaiKeycheck.check_key(
            key, zaiKeycheck.base_urls_from_environment(), proxy, timeout, debug,
        )
    if provider == "qwen":
        from keycheckers.qwen import qwenKeycheck

        custom = qwenKeycheck.split_csv(
            os.getenv("QWEN_BASE_URLS") or os.getenv("DASHSCOPE_BASE_URLS")
        )
        base_urls = qwenKeycheck.unique_ordered([*custom, *qwenKeycheck.DEFAULT_BASE_URLS])
        return qwenKeycheck.check_key(key, base_urls, bool(custom), proxy, timeout, debug)
    if provider == "kimi":
        from keycheckers.kimi import kimiKeycheck

        custom = kimiKeycheck.split_csv(
            os.getenv("KIMI_BASE_URLS") or os.getenv("MOONSHOT_BASE_URLS")
        )
        base_urls = kimiKeycheck.unique_ordered([*custom, *kimiKeycheck.DEFAULT_BASE_URLS])
        return kimiKeycheck.check_key(key, base_urls, proxy, timeout, debug)
    raise ValueError(f"unsupported provider resolution adapter: {provider}")


def resolve_provider_key(
    key, finding=None, proxy=None, timeout=15, debug=False, hint="",
    origin_service="", configured_order=None, probe=None,
):
    providers = ordered_providers(finding, hint, origin_service, configured_order)
    probe = probe or default_provider_probe
    attempts = []
    retry_results = []
    for provider in providers:
        result = probe(provider, key, proxy, timeout, debug)
        result = result if isinstance(result, dict) else {"status": "UNKNOWN"}
        outcome = provider_result_outcome(result)
        attempts.append(bounded_attempt(provider, result, outcome))
        if outcome == "match":
            return {
                **result,
                "resolved_provider": provider,
                "provider_resolution": "matched",
                "provider_resolution_order": providers,
                "provider_resolution_attempts": attempts,
                "result_source": "provider_resolution",
            }
        if outcome == "retry":
            retry_results.append(result)

    if retry_results:
        selected = retry_results[0]
        return {
            **selected,
            "provider_resolution": "retry",
            "provider_resolution_order": providers,
            "provider_resolution_attempts": attempts,
            "result_source": "provider_resolution",
            "message": str(selected.get("message") or "provider resolution remains inconclusive")[:1000],
        }
    return {
        "status": "DEAD",
        "provider_resolution": "exhausted",
        "provider_resolution_order": providers,
        "provider_resolution_attempts": attempts,
        "result_source": "provider_resolution",
        "message": "all compatible providers rejected the credential",
    }