Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,189 @@
|
||||
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",
|
||||
}
|
||||
Reference in New Issue
Block a user