truf-server/app/keycheckers/zai/zaiKeycheck.py

import sys

sys.dont_write_bytecode = True

import argparse
import json
import os
import re
from urllib.parse import urlparse

import requests

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from keycheck_common import (
    combined_provider_routing_hint,
    commit_status_transaction,
    default_input_file,
    default_proxy_file,
    ensure_output_files,
    iter_findings,
    keycheck_input_mode,
    load_checked_statuses,
    load_known_keys,
    load_proxies,
    mask_secret,
    provider_routing_database_failed,
    read_plain_keys,
    record_validation_result,
    recover_status_transaction,
    require_provider_authority,
    service_output_dir,
    should_skip_key,
    write_keycheck_event,
)
from keycheckers.provider_resolution import (
    AMBIGUOUS_GENERIC_SK_HINT,
    AMBIGUOUS_QWEN_DEEPSEEK_HINT,
    resolve_provider_key,
)


SERVICE = "zai"
DETECTOR = "ZaiGLM"
OUTPUT_DIR = os.getenv("KEYCHECK_OUTPUT_DIR") or service_output_dir(SERVICE)
INPUT_FILE = os.getenv("KEYCHECK_INPUT_FILE") or default_input_file()
PROXY_FILE = os.getenv("KEYCHECK_PROXY_FILE") or default_proxy_file()

CHECKED_FILE = os.path.join(OUTPUT_DIR, "zaiChecked.txt")
RESULTS_FILE = os.path.join(OUTPUT_DIR, "zaiResults.jsonl")
STATUS_FILES = {
    "VALID": os.path.join(OUTPUT_DIR, "zaiAlive.txt"),
    "NO_BALANCE": os.path.join(OUTPUT_DIR, "zaiNoBalance.txt"),
    "DEAD": os.path.join(OUTPUT_DIR, "zaiDead.txt"),
    "RESTRICTED": os.path.join(OUTPUT_DIR, "zaiRestricted.txt"),
    "LIMITED": os.path.join(OUTPUT_DIR, "zaiLimited.txt"),
    "NETWORK": os.path.join(OUTPUT_DIR, "zaiNetwork.txt"),
    "UNKNOWN": os.path.join(OUTPUT_DIR, "zaiUnknown.txt"),
}

DEFAULT_BASE_URLS = (
    "https://api.z.ai/api/paas/v4",
    "https://open.bigmodel.cn/api/paas/v4",
)
ZAI_KEY_REGEX = re.compile(
    r"(?<![A-Za-z0-9_.-])(?:"
    r"(?:zai|sk)-[A-Za-z0-9][A-Za-z0-9_-]{20,505}|"
    r"[A-Fa-f0-9]{32}\.[A-Za-z0-9_-]{16,128}"
    r")(?![A-Za-z0-9_.-])"
)
ZAI_DOTTED_KEY_REGEX = re.compile(r"^[A-Fa-f0-9]{32}\.[A-Za-z0-9_-]{16,128}$")
ZAI_CONTEXT_REGEX = re.compile(
    r"(?:ZAI_API_KEY|GLM_API_KEY|ZHIPUAI_API_KEY|BIGMODEL_API_KEY|api\.z\.ai|"
    r"open\.bigmodel\.cn|zhipuai|chatglm)",
    re.IGNORECASE,
)
FOREIGN_KEY_PREFIXES = ("sk-ant-", "sk-or-", "sk-proj-", "sk-svcacct-", "sk-admin-")
AMBIGUOUS_HINTS = {AMBIGUOUS_QWEN_DEEPSEEK_HINT, AMBIGUOUS_GENERIC_SK_HINT}
AUTH_FAILURE_CODES = {"1000", "1001", "1003"}
AUTHENTICATED_NO_BALANCE_CODES = {"1113"}
AUTHENTICATED_LIMIT_CODES = {"1302", "1308", "1309", "1310", "1311"}
AUTHENTICATED_RESTRICTED_CODES = {"1005", "1220"}
PROBE_MODEL = "glm-5.2"


def normalize_base_url(value):
    return str(value or "").strip().rstrip("/")


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


def unique_ordered(values):
    output = []
    seen = set()
    for value in values:
        normalized = normalize_base_url(value)
        if normalized and normalized not in seen:
            seen.add(normalized)
            output.append(normalized)
    return output


def base_urls_from_environment(extra=None, include_defaults=True):
    configured = []
    for value in extra or ():
        configured.extend(split_csv(value))
    configured.extend(split_csv(os.getenv("ZAI_BASE_URLS") or os.getenv("ZHIPU_BASE_URLS")))
    defaults = DEFAULT_BASE_URLS if include_defaults else ()
    return unique_ordered([*configured, *defaults])


def endpoint_label(base_url):
    parsed = urlparse(base_url)
    return parsed.netloc or base_url


def key_from_text(*values):
    for value in values:
        match = ZAI_KEY_REGEX.search(str(value or ""))
        if match:
            return match.group(0)
    return ""


def key_rejection_reason(key):
    value = str(key or "")
    try:
        encoded = value.encode("utf-8", errors="strict")
    except UnicodeEncodeError:
        return "candidate is not valid UTF-8"
    if len(encoded) > 512:
        return "candidate exceeds the 512-byte key limit"
    if value.startswith(FOREIGN_KEY_PREFIXES):
        return "candidate has a foreign provider prefix"
    if not ZAI_KEY_REGEX.fullmatch(value):
        return "candidate does not match a bounded ZAI key format"
    return ""


def finding_detector_names(finding):
    if not isinstance(finding, dict):
        return set()
    extra = finding.get("ExtraData") if isinstance(finding.get("ExtraData"), dict) else {}
    return {
        name for name in (
            str(finding.get("DetectorName") or finding.get("DetectorType") or "").strip().lower(),
            str(extra.get("name") or "").strip().lower(),
        ) if name
    }


def finding_provider_routing_hint(finding, key=""):
    if not isinstance(finding, dict):
        return "zai" if ZAI_DOTTED_KEY_REGEX.fullmatch(str(key or "")) else ""
    context = finding.get("ScannerContext") if isinstance(finding.get("ScannerContext"), dict) else {}
    persisted = str(context.get("provider_hint") or "").strip().lower()
    if persisted:
        return persisted
    if "zaiglm" in finding_detector_names(finding) or ZAI_DOTTED_KEY_REGEX.fullmatch(str(key or "")):
        return "zai"
    text = "\n".join(str(context.get(name) or "") for name in ("nearby", "file"))
    return "zai" if ZAI_CONTEXT_REGEX.search(text) else ""


def iter_candidate_decisions(input_file, plain_files, trusted_retry_files=None):
    detectors = [
        "ZaiGLM", "zaiglm", "CustomRegex", "QwenDashScope", "Qwen_DashScope",
        "Qwen", "DashScope", "DeepSeek", "DeepSeekApiKey", "DeepSeek_API_Key",
        "KimiMoonshot", "MoonshotAI", "Moonshot", "Kimi",
    ]
    seen = set()
    for item in iter_findings(input_file, detectors):
        finding = item.get("finding") or {}
        key = item.get("credential_secret_text") or key_from_text(
            item.get("raw"), item.get("raw_v2"), finding.get("Raw"), finding.get("RawV2"),
        )
        if not key or key_rejection_reason(key):
            continue
        local_hint = finding_provider_routing_hint(finding, key)
        hint = combined_provider_routing_hint(key, local_hint)
        if provider_routing_database_failed():
            raise RuntimeError("provider routing evidence lookup failed closed")
        seen.add(key)
        yield key, item.get("source") or input_file, finding, hint

    owned_retry_paths = {
        os.path.normcase(os.path.abspath(path)) for path in STATUS_FILES.values()
    }
    retry_paths = [
        path for path in trusted_retry_files or ()
        if os.path.normcase(os.path.abspath(path)) in owned_retry_paths
    ]
    for item in read_plain_keys([*plain_files, *retry_paths], ZAI_KEY_REGEX):
        key = item["key"]
        if key in seen or key_rejection_reason(key):
            continue
        yield key, item["source"], {}, "zai"


def redact_text(value, key):
    text = str(value or "")[:1000]
    if key:
        text = text.replace(key, "***REDACTED***")
    return ZAI_KEY_REGEX.sub("***REDACTED***", text)


def parse_error(response, key):
    try:
        payload = response.json()
    except ValueError:
        payload = {}
    error = payload.get("error") if isinstance(payload, dict) else {}
    if not isinstance(error, dict):
        error = {}
    return {
        "http_status": int(response.status_code),
        "code": str(error.get("code") or (payload.get("code") if isinstance(payload, dict) else "") or ""),
        "message": redact_text(
            error.get("message") or error.get("msg") or (
                payload.get("message") or payload.get("msg") if isinstance(payload, dict) else ""
            ) or response.text[:500],
            key,
        ),
    }


def classify_error(error):
    http_status = int(error.get("http_status") or 0)
    code = str(error.get("code") or "")
    message = str(error.get("message") or "").lower()
    if http_status == 402 or any(marker in message for marker in (
        "insufficient balance", "balance is insufficient", "no balance", "account balance",
        "recharge", "payment required", "billing arrears", "credit balance",
    )):
        return "NO_BALANCE", True
    if code in AUTHENTICATED_NO_BALANCE_CODES:
        return "NO_BALANCE", True
    if any(marker in message for marker in (
        "quota", "rate limit", "rate-limit", "too many requests", "resource exhausted",
        "concurrency limit", "usage limit",
    )):
        return "LIMITED", True
    if code in AUTHENTICATED_LIMIT_CODES:
        return "LIMITED", True
    if any(marker in message for marker in (
        "permission denied", "access denied", "not authorized for", "model access", "forbidden",
    )):
        return "RESTRICTED", True
    if code in AUTHENTICATED_RESTRICTED_CODES:
        return "RESTRICTED", True
    if http_status == 401 or code in AUTH_FAILURE_CODES:
        return "DEAD", False
    if 500 <= http_status <= 599 or code in {"1200", "1230", "1234", "1305"}:
        return "NETWORK", False
    if http_status == 429:
        return "LIMITED", False
    if http_status == 403:
        return "RESTRICTED", False
    return "UNKNOWN", False


def probe_chat_completion(key, base_url, model, proxy, timeout, debug=False):
    if not model:
        return {
            "status": "UNKNOWN", "model": "",
            "message": "no chat-capable model returned by /models",
        }
    url = f"{normalize_base_url(base_url)}/chat/completions"
    headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "ping"}],
        "max_tokens": 1,
        "stream": False,
    }
    try:
        response = requests.post(url, headers=headers, json=payload, proxies=proxy, timeout=timeout)
    except requests.RequestException as exc:
        return {
            "status": "NETWORK", "model": model,
            "message": redact_text(exc, key),
        }
    if debug:
        print(
            f"  DEBUG {endpoint_label(base_url)} chat probe {model}: HTTP {response.status_code}: "
            f"{redact_text(response.text[:500], key)}"
        )
    if response.status_code == 200:
        try:
            response_payload = response.json()
        except ValueError as exc:
            return {
                "status": "UNKNOWN", "model": model, "http_status": response.status_code,
                "message": f"invalid chat completion response: {exc}",
            }
        if isinstance(response_payload, dict) and response_payload.get("choices"):
            return {
                "status": "GENERATION_OK", "model": model,
                "http_status": response.status_code, "message": "chat completion accepted",
            }
    error = parse_error(response, key)
    status, authenticated = classify_error(error)
    return {
        "status": status, "authenticated": authenticated, "model": model,
        "http_status": response.status_code, "business_code": error.get("code") or "",
        "error": error, "message": error.get("message") or "",
    }


def check_base_url(key, base_url, proxy, timeout, debug=False):
    url = f"{normalize_base_url(base_url)}/models"
    headers = {"Authorization": f"Bearer {key}", "Accept": "application/json"}
    try:
        response = requests.get(url, headers=headers, proxies=proxy, timeout=timeout)
    except requests.RequestException as exc:
        return {
            "status": "NETWORK", "base_url": base_url,
            "region": endpoint_label(base_url), "message": redact_text(exc, key),
        }
    if debug:
        print(
            f"  DEBUG {endpoint_label(base_url)} /models: HTTP {response.status_code}: "
            f"{redact_text(response.text[:500], key)}"
        )
    if response.status_code == 200:
        try:
            payload = response.json()
            data = payload.get("data") if isinstance(payload, dict) else None
            if not isinstance(data, list):
                raise ValueError("missing data model list")
            models = sorted({
                str(item.get("id") or item.get("name") or "")
                for item in data if isinstance(item, dict) and (item.get("id") or item.get("name"))
            })
        except (TypeError, ValueError, json.JSONDecodeError) as exc:
            return {
                "status": "UNKNOWN", "base_url": base_url,
                "region": endpoint_label(base_url), "message": f"invalid models response: {exc}",
            }
        probe_model = PROBE_MODEL
        probe = probe_chat_completion(key, base_url, probe_model, proxy, timeout, debug)
        probe_status = probe.get("status") or "UNKNOWN"
        if probe_status == "GENERATION_OK":
            status = "VALID"
        elif probe_status in {"NO_BALANCE", "LIMITED", "RESTRICTED", "NETWORK"}:
            status = probe_status
        elif probe_status == "DEAD":
            status = "RESTRICTED"
        else:
            status = "UNKNOWN"
        probe_message = probe.get("message") or json.dumps(probe.get("error") or {}, ensure_ascii=False)
        message = (
            f"chat probe ok; model={probe_model}; models={len(data)}"
            if status == "VALID"
            else f"models authenticated; generation_probe={probe_status}; model={probe_model}; "
                 f"models={len(data)}; {probe_message}"
        )
        return {
            "status": status, "authenticated": True, "base_url": base_url,
            "region": endpoint_label(base_url), "model_count": len(data),
            "models": models[:30], "llm_probe_status": probe_status,
            "llm_probe_model": probe.get("model") or probe_model,
            "llm_probe_http_status": probe.get("http_status"),
            "business_code": probe.get("business_code") or "",
            "probe": probe, "error": probe.get("error") or {},
            "message": message[:1000],
        }
    error = parse_error(response, key)
    status, authenticated = classify_error(error)
    return {
        "status": status, "authenticated": authenticated, "base_url": base_url,
        "region": endpoint_label(base_url), "http_status": response.status_code,
        "business_code": error.get("code") or "", "error": error,
        "message": error.get("message") or "",
    }


def check_key(key, base_urls, proxy, timeout, debug=False):
    rejection = key_rejection_reason(key)
    if rejection:
        return {"status": "UNKNOWN", "message": rejection, "candidate_rejected": True}
    attempts = []
    for base_url in base_urls:
        result = check_base_url(key, base_url, proxy, timeout, debug)
        attempts.append(result)
        if result.get("authenticated"):
            return {**result, "attempts": attempts}
    statuses = [attempt.get("status") for attempt in attempts]
    status = next(
        (candidate for candidate in ("NETWORK", "LIMITED", "RESTRICTED", "UNKNOWN", "DEAD") if candidate in statuses),
        "UNKNOWN",
    )
    selected = next((attempt for attempt in attempts if attempt.get("status") == status), {})
    return {**selected, "status": status, "attempts": attempts}


def ensure_files():
    ensure_output_files([CHECKED_FILE, RESULTS_FILE, *STATUS_FILES.values()])
    recover_status_transaction(CHECKED_FILE, STATUS_FILES)


def write_result(key, result, source, finding):
    status = result.get("status") or "UNKNOWN"
    write_keycheck_event(SERVICE, RESULTS_FILE, key, result, source, finding, DETECTOR)
    commit_status_transaction(
        CHECKED_FILE, STATUS_FILES, key, status,
        result.get("message", ""), result.get("region") or source,
    )
    record_validation_result(SERVICE, key, result, source, finding, DETECTOR)


def retry_statuses_from_args(args):
    statuses = set()
    for enabled, status in (
        (args.retry_network, "NETWORK"),
        (args.retry_limited, "LIMITED"),
        (args.retry_unknown, "UNKNOWN"),
        (args.retry_restricted, "RESTRICTED"),
        (args.retry_no_balance, "NO_BALANCE"),
        (args.retry_valid, "VALID"),
    ):
        if enabled:
            statuses.add(status)
    return statuses


def retry_input_files_from_args(args):
    if args.recheck_all:
        statuses = list(STATUS_FILES)
    else:
        statuses = list(retry_statuses_from_args(args))
    return [STATUS_FILES[status] for status in statuses]


def parse_args():
    parser = argparse.ArgumentParser(description="ZAI / Zhipu GLM key checker")
    parser.add_argument("--input", default=INPUT_FILE)
    parser.add_argument("--plain", action="append", default=[])
    parser.add_argument("--proxy-file", default=PROXY_FILE)
    parser.add_argument("--timeout", type=int, default=15)
    parser.add_argument("--max-keys", type=int, default=0)
    parser.add_argument("--base-url", action="append", default=[])
    parser.add_argument("--no-default-base-urls", action="store_true")
    parser.add_argument("--retry-network", action="store_true")
    parser.add_argument("--retry-limited", action="store_true")
    parser.add_argument("--retry-unknown", action="store_true")
    parser.add_argument("--retry-restricted", action="store_true")
    parser.add_argument("--retry-no-balance", action="store_true")
    parser.add_argument("--retry-valid", action="store_true")
    parser.add_argument("--recheck-all", action="store_true")
    parser.add_argument("--debug", action="store_true")
    return parser.parse_args()


def main():
    require_provider_authority(SERVICE)
    args = parse_args()
    ensure_files()
    proxy_cycler = load_proxies(args.proxy_file)
    checked = load_checked_statuses(CHECKED_FILE)
    known = load_known_keys(CHECKED_FILE, STATUS_FILES)
    retry_statuses = retry_statuses_from_args(args)
    retry_files = retry_input_files_from_args(args)
    base_urls = base_urls_from_environment(args.base_url, not args.no_default_base_urls)
    if not base_urls:
        raise SystemExit("No ZAI base URLs configured")

    processed = 0
    skipped = 0
    postgres_mode = keycheck_input_mode() == "postgres"
    for key, source, finding, routing_hint in iter_candidate_decisions(args.input, args.plain, retry_files):
        is_ambiguous = routing_hint in AMBIGUOUS_HINTS
        if routing_hint != "zai" and not (postgres_mode and is_ambiguous):
            skipped += 1
            continue
        if should_skip_key(
            key, checked, known, args, retry_statuses, service=SERVICE,
            source=source, finding=finding, detector=DETECTOR,
        ):
            skipped += 1
            continue
        if args.max_keys and processed >= args.max_keys:
            break
        processed += 1
        print(f"\n[{processed}] ZAI candidate {mask_secret(key)} from {source}")
        proxy = next(proxy_cycler) if proxy_cycler else None
        if is_ambiguous:
            result = resolve_provider_key(
                key, finding, proxy, args.timeout, args.debug,
                hint=routing_hint, origin_service=SERVICE,
            )
        else:
            result = check_key(key, base_urls, proxy, args.timeout, args.debug)
        print(f"  STATUS: {result['status']} | {result.get('message', '')[:200]}")
        write_result(key, result, source, finding)
        known.add(key)
        checked[key] = result["status"]

    print(f"\nDone. Processed={processed}, skipped={skipped}, results={RESULTS_FILE}")


if __name__ == "__main__":
    main()