Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,231 @@
|
||||
import sys
|
||||
|
||||
sys.dont_write_bytecode = True
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
|
||||
import requests
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from keycheck_common import (
|
||||
append_jsonl, classify_common_http_status, 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,
|
||||
read_plain_keys, record_validation_result, recover_status_transaction,
|
||||
request_error_message, require_provider_authority, service_output_dir, should_skip_key, write_keycheck_event,
|
||||
)
|
||||
|
||||
SERVICE = "xai"
|
||||
DETECTOR_NAMES = ["XAI", "XAi", "Xai"]
|
||||
DETECTOR = "XAI"
|
||||
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, "xaiChecked.txt")
|
||||
RESULTS_FILE = os.path.join(OUTPUT_DIR, "xaiResults.jsonl")
|
||||
STATUS_FILES = {
|
||||
"VALID": os.path.join(OUTPUT_DIR, "xaiAlive.txt"),
|
||||
"NO_BALANCE": os.path.join(OUTPUT_DIR, "xaiNoBalance.txt"),
|
||||
"DEAD": os.path.join(OUTPUT_DIR, "xaiDead.txt"),
|
||||
"RESTRICTED": os.path.join(OUTPUT_DIR, "xaiRestricted.txt"),
|
||||
"LIMITED": os.path.join(OUTPUT_DIR, "xaiLimited.txt"),
|
||||
"NO_CONTEXT": os.path.join(OUTPUT_DIR, "xaiNoContext.txt"),
|
||||
"NETWORK": os.path.join(OUTPUT_DIR, "xaiNetwork.txt"),
|
||||
"UNKNOWN": os.path.join(OUTPUT_DIR, "xaiUnknown.txt"),
|
||||
}
|
||||
KEY_REGEX = re.compile(r"\bxai-[A-Za-z0-9_-]{20,}\b")
|
||||
MODELS_URL = "https://api.x.ai/v1/models"
|
||||
CHAT_URL = "https://api.x.ai/v1/chat/completions"
|
||||
CHAT_MODEL_PRIORITY = (
|
||||
"grok-4.6",
|
||||
"grok-4.5",
|
||||
"grok-4.3",
|
||||
"grok-4.20-0309-reasoning",
|
||||
"grok-4.20-0309-non-reasoning",
|
||||
)
|
||||
NO_BALANCE_MARKERS = (
|
||||
"quota",
|
||||
"billing",
|
||||
"balance",
|
||||
"credit",
|
||||
"credits",
|
||||
"payment",
|
||||
"insufficient",
|
||||
"depleted",
|
||||
"spending limit",
|
||||
"no credits",
|
||||
"used all available credits",
|
||||
"doesn't have any credits",
|
||||
)
|
||||
DEAD_MARKERS = ("incorrect api key", "invalid api key", "api key provided", "invalid-argument")
|
||||
|
||||
|
||||
def ensure_files():
|
||||
ensure_output_files([CHECKED_FILE, RESULTS_FILE, *STATUS_FILES.values()])
|
||||
recover_status_transaction(CHECKED_FILE, STATUS_FILES)
|
||||
|
||||
|
||||
def iter_candidate_decisions(input_file, plain_files):
|
||||
seen_plain = set()
|
||||
for item in iter_findings(input_file, DETECTOR_NAMES):
|
||||
key = item.get("credential_secret_text") or item["raw"]
|
||||
if key:
|
||||
yield key, item["source"], item["finding"], bool(KEY_REGEX.fullmatch(key))
|
||||
for item in read_plain_keys(plain_files, KEY_REGEX):
|
||||
key = item["key"]
|
||||
if key not in seen_plain:
|
||||
seen_plain.add(key)
|
||||
yield key, item["source"], {}, True
|
||||
|
||||
|
||||
def extract_candidates(input_file, plain_files):
|
||||
for key, source, finding, valid_format in iter_candidate_decisions(input_file, plain_files):
|
||||
if valid_format:
|
||||
yield key, source, finding
|
||||
|
||||
|
||||
def check_key(key, proxy, timeout):
|
||||
try:
|
||||
response = requests.get(MODELS_URL, headers={"Authorization": f"Bearer {key}", "Accept": "application/json"}, proxies=proxy, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
return {"status": "NETWORK", "message": str(exc)}
|
||||
if response.status_code == 200:
|
||||
data = response.json() if response.text else {}
|
||||
models = [item.get("id") for item in data.get("data", []) if isinstance(item, dict) and item.get("id")]
|
||||
model_inventory = sorted(set(models))
|
||||
model = choose_chat_model(models)
|
||||
probe = probe_chat_completion(key, model, proxy, timeout)
|
||||
if probe.get("status") != "GENERATION_OK":
|
||||
return {
|
||||
"status": probe.get("status") or "UNKNOWN",
|
||||
"message": probe.get("message", ""),
|
||||
"model_count": len(models),
|
||||
"models": models[:20],
|
||||
"model_inventory": model_inventory,
|
||||
"llm_probe_status": probe.get("status"),
|
||||
"llm_probe_model": probe.get("model", model),
|
||||
"llm_probe_http_status": probe.get("http_status"),
|
||||
}
|
||||
return {
|
||||
"status": "VALID", "message": f"chat ping ok; model={model}; models={len(models)}",
|
||||
"model_count": len(models), "models": models[:20], "model_inventory": model_inventory,
|
||||
"llm_probe_status": probe.get("status"), "llm_probe_model": model,
|
||||
}
|
||||
status = classify_xai_response(response)
|
||||
return {"status": status, "http_status": response.status_code, "message": request_error_message(response).replace(key, "***REDACTED***")}
|
||||
|
||||
|
||||
def classify_xai_response(response):
|
||||
message = request_error_message(response).lower()
|
||||
if any(marker in message for marker in DEAD_MARKERS):
|
||||
return "DEAD"
|
||||
if any(marker in message for marker in NO_BALANCE_MARKERS):
|
||||
return "NO_BALANCE"
|
||||
if response.status_code == 401:
|
||||
return "DEAD"
|
||||
if response.status_code == 403:
|
||||
return "RESTRICTED"
|
||||
if response.status_code == 429:
|
||||
return "LIMITED"
|
||||
return classify_common_http_status(response.status_code)
|
||||
|
||||
|
||||
def choose_chat_model(models):
|
||||
models = [str(model or "") for model in models if model]
|
||||
by_lower = {model.lower(): model for model in models}
|
||||
for model in CHAT_MODEL_PRIORITY:
|
||||
if model.lower() in by_lower:
|
||||
return by_lower[model.lower()]
|
||||
for model in models:
|
||||
if "grok" in model.lower():
|
||||
return model
|
||||
return models[0] if models else ""
|
||||
|
||||
|
||||
def probe_chat_completion(key, model, proxy, timeout):
|
||||
if not model:
|
||||
return {"status": "NO_CONTEXT", "message": "no chat-capable model from /models", "model": ""}
|
||||
headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
|
||||
payload = {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 1}
|
||||
try:
|
||||
response = requests.post(CHAT_URL, headers=headers, json=payload, proxies=proxy, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
return {"status": "NETWORK", "message": str(exc), "model": model}
|
||||
if response.status_code == 200:
|
||||
return {"status": "GENERATION_OK", "message": "chat completion accepted", "model": model}
|
||||
return {"status": classify_xai_response(response), "http_status": response.status_code, "message": request_error_message(response).replace(key, "***REDACTED***"), "model": model}
|
||||
|
||||
|
||||
def write_result(key, result, source, finding):
|
||||
status = result.get("status") or "UNKNOWN"
|
||||
write_keycheck_event(SERVICE, RESULTS_FILE, key, result, source, finding, DETECTOR)
|
||||
extra = ",".join(result.get("models") or [])[:500] if status == "VALID" else source
|
||||
commit_status_transaction(
|
||||
CHECKED_FILE, STATUS_FILES, key, status, result.get("message", ""), extra,
|
||||
)
|
||||
record_validation_result(SERVICE, key, result, source, finding, DETECTOR)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="xAI 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("--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")
|
||||
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 = set()
|
||||
if args.retry_network: retry_statuses.add("NETWORK")
|
||||
if args.retry_limited: retry_statuses.add("LIMITED")
|
||||
if args.retry_unknown: retry_statuses.update({"UNKNOWN", "NO_CONTEXT"})
|
||||
if args.retry_restricted: retry_statuses.add("RESTRICTED")
|
||||
if args.retry_no_balance: retry_statuses.add("NO_BALANCE")
|
||||
if args.retry_valid: retry_statuses.add("VALID")
|
||||
processed = skipped = 0
|
||||
print("--- xAI key checker ---")
|
||||
postgres_mode = keycheck_input_mode() == "postgres"
|
||||
for key, source, finding, valid_format in iter_candidate_decisions(args.input, args.plain):
|
||||
if not valid_format and not postgres_mode:
|
||||
skipped += 1
|
||||
continue
|
||||
if valid_format and 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}] xAI candidate {mask_secret(key)} from {source}")
|
||||
result = (
|
||||
check_key(key, next(proxy_cycler) if proxy_cycler else None, args.timeout)
|
||||
if valid_format else
|
||||
{"status": "NO_CONTEXT", "message": "candidate does not match canonical xAI token format"}
|
||||
)
|
||||
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()
|
||||
Reference in New Issue
Block a user