Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,403 @@
|
||||
import sys
|
||||
|
||||
sys.dont_write_bytecode = True
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
from collections import Counter
|
||||
|
||||
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 = "replicate"
|
||||
DETECTOR = "Replicate"
|
||||
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, "replicateChecked.txt")
|
||||
RESULTS_FILE = os.path.join(OUTPUT_DIR, "replicateResults.jsonl")
|
||||
STATUS_FILES = {
|
||||
"VALID": os.path.join(OUTPUT_DIR, "replicateAlive.txt"),
|
||||
"DEAD": os.path.join(OUTPUT_DIR, "replicateDead.txt"),
|
||||
"RESTRICTED": os.path.join(OUTPUT_DIR, "replicateRestricted.txt"),
|
||||
"LIMITED": os.path.join(OUTPUT_DIR, "replicateLimited.txt"),
|
||||
"NO_BALANCE": os.path.join(OUTPUT_DIR, "replicateNoBalance.txt"),
|
||||
"NETWORK": os.path.join(OUTPUT_DIR, "replicateNetwork.txt"),
|
||||
"NO_CONTEXT": os.path.join(OUTPUT_DIR, "replicateNoContext.txt"),
|
||||
"UNKNOWN": os.path.join(OUTPUT_DIR, "replicateUnknown.txt"),
|
||||
}
|
||||
KEY_REGEX = re.compile(r"\br8_[A-Za-z0-9]{30,}\b")
|
||||
API_BASE = "https://api.replicate.com/v1"
|
||||
ACCOUNT_URL = f"{API_BASE}/account"
|
||||
RESOURCE_ENDPOINTS = {
|
||||
"predictions": f"{API_BASE}/predictions",
|
||||
"deployments": f"{API_BASE}/deployments",
|
||||
"trainings": f"{API_BASE}/trainings",
|
||||
}
|
||||
NO_BALANCE_MARKERS = (
|
||||
"balance",
|
||||
"billing",
|
||||
"credit",
|
||||
"credits",
|
||||
"payment",
|
||||
"insufficient",
|
||||
"depleted",
|
||||
"no credits",
|
||||
"out of credit",
|
||||
"run out of credit",
|
||||
)
|
||||
|
||||
|
||||
def ensure_files():
|
||||
ensure_output_files([CHECKED_FILE, RESULTS_FILE, *STATUS_FILES.values()])
|
||||
recover_status_transaction(CHECKED_FILE, STATUS_FILES)
|
||||
|
||||
|
||||
def auth_headers(key):
|
||||
return {"Authorization": f"Bearer {key}", "Accept": "application/json"}
|
||||
|
||||
|
||||
def redacted_error_message(response, key):
|
||||
return request_error_message(response).replace(key, "***REDACTED***")
|
||||
|
||||
|
||||
def classify_replicate_response(response, key):
|
||||
message = redacted_error_message(response, key).lower()
|
||||
if response.status_code == 402 or any(marker in message for marker in NO_BALANCE_MARKERS):
|
||||
return "NO_BALANCE"
|
||||
if response.status_code == 403:
|
||||
return "RESTRICTED"
|
||||
return classify_common_http_status(response.status_code)
|
||||
|
||||
|
||||
def api_get(key, url, proxy, timeout, debug=False):
|
||||
try:
|
||||
response = requests.get(url, headers=auth_headers(key), proxies=proxy, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
return {"status": "NETWORK", "message": str(exc)[:1000], "payload": None}
|
||||
if debug:
|
||||
detail = "ok" if response.status_code == 200 else redacted_error_message(response, key)[:500]
|
||||
print(f" DEBUG GET {url}: HTTP {response.status_code}: {detail}")
|
||||
if response.status_code != 200:
|
||||
return {
|
||||
"status": classify_replicate_response(response, key),
|
||||
"http_status": response.status_code,
|
||||
"message": redacted_error_message(response, key),
|
||||
"payload": None,
|
||||
}
|
||||
try:
|
||||
payload = response.json() if response.text else {}
|
||||
except ValueError:
|
||||
payload = {}
|
||||
return {"status": "OK", "http_status": 200, "message": "ok", "payload": payload}
|
||||
|
||||
|
||||
def paginated_items(payload):
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
if not isinstance(payload, dict):
|
||||
return []
|
||||
for key in ("results", "data", "items"):
|
||||
value = payload.get(key)
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
return []
|
||||
|
||||
|
||||
def text_value(value):
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def compact_model_ref(value):
|
||||
if isinstance(value, str):
|
||||
return value.strip()
|
||||
if not isinstance(value, dict):
|
||||
return ""
|
||||
owner = text_value(value.get("owner") or value.get("model_owner"))
|
||||
name = text_value(value.get("name") or value.get("model_name"))
|
||||
if owner and name:
|
||||
return f"{owner}/{name}"
|
||||
for key in ("model", "id", "slug"):
|
||||
item = text_value(value.get(key))
|
||||
if item:
|
||||
return item
|
||||
url = text_value(value.get("url") or value.get("web_url"))
|
||||
if "replicate.com/" in url:
|
||||
return url.rstrip("/").split("replicate.com/", 1)[-1]
|
||||
return ""
|
||||
|
||||
|
||||
def model_refs_from_item(item):
|
||||
if not isinstance(item, dict):
|
||||
return []
|
||||
refs = []
|
||||
for key in ("model", "destination", "source_model", "base_model"):
|
||||
ref = compact_model_ref(item.get(key))
|
||||
if ref:
|
||||
refs.append(ref)
|
||||
for key in ("version", "latest_version", "current_release"):
|
||||
value = item.get(key)
|
||||
if isinstance(value, dict):
|
||||
ref = compact_model_ref(value.get("model") or value.get("destination"))
|
||||
if ref:
|
||||
refs.append(ref)
|
||||
return sorted(set(refs))
|
||||
|
||||
|
||||
def summarize_predictions(payload, limit=10):
|
||||
items = paginated_items(payload)
|
||||
models = sorted({text_value(item.get("model")) for item in items if isinstance(item, dict) and item.get("model")})
|
||||
statuses = Counter(text_value(item.get("status")) for item in items if isinstance(item, dict) and item.get("status"))
|
||||
samples = []
|
||||
for item in items[:limit]:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
samples.append({
|
||||
"id": text_value(item.get("id"))[:80],
|
||||
"status": text_value(item.get("status")),
|
||||
"model": text_value(item.get("model")),
|
||||
"source": text_value(item.get("source")),
|
||||
"data_removed": bool(item.get("data_removed")),
|
||||
"created_at": text_value(item.get("created_at")),
|
||||
"completed_at": text_value(item.get("completed_at")),
|
||||
})
|
||||
return {
|
||||
"prediction_count_sample": len(items),
|
||||
"prediction_has_next_page": bool(isinstance(payload, dict) and payload.get("next")),
|
||||
"prediction_status_counts": dict(statuses),
|
||||
"prediction_models": models[:50],
|
||||
"prediction_samples": samples,
|
||||
}
|
||||
|
||||
|
||||
def deployment_name(item):
|
||||
owner = text_value(item.get("owner") or item.get("deployment_owner"))
|
||||
name = text_value(item.get("name") or item.get("deployment_name"))
|
||||
if owner and name and "/" not in name:
|
||||
return f"{owner}/{name}"
|
||||
return name or owner
|
||||
|
||||
|
||||
def summarize_deployments(payload, limit=20):
|
||||
items = paginated_items(payload)
|
||||
models = sorted({ref for item in items for ref in model_refs_from_item(item)})
|
||||
deployments = []
|
||||
for item in items[:limit]:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
current_release = item.get("current_release") if isinstance(item.get("current_release"), dict) else {}
|
||||
deployments.append({
|
||||
"name": deployment_name(item),
|
||||
"model": next(iter(model_refs_from_item(item)), ""),
|
||||
"version": text_value(item.get("version") or current_release.get("version"))[:80],
|
||||
"hardware": text_value(item.get("hardware") or current_release.get("hardware")),
|
||||
"min_instances": item.get("min_instances"),
|
||||
"max_instances": item.get("max_instances"),
|
||||
})
|
||||
return {
|
||||
"deployment_count": len(items),
|
||||
"deployment_has_next_page": bool(isinstance(payload, dict) and payload.get("next")),
|
||||
"deployment_models": models[:50],
|
||||
"deployments": deployments,
|
||||
}
|
||||
|
||||
|
||||
def summarize_trainings(payload, limit=10):
|
||||
items = paginated_items(payload)
|
||||
models = sorted({ref for item in items for ref in model_refs_from_item(item)})
|
||||
statuses = Counter(text_value(item.get("status")) for item in items if isinstance(item, dict) and item.get("status"))
|
||||
samples = []
|
||||
for item in items[:limit]:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
samples.append({
|
||||
"id": text_value(item.get("id"))[:80],
|
||||
"status": text_value(item.get("status")),
|
||||
"model": next(iter(model_refs_from_item(item)), ""),
|
||||
"created_at": text_value(item.get("created_at")),
|
||||
"completed_at": text_value(item.get("completed_at")),
|
||||
})
|
||||
return {
|
||||
"training_count_sample": len(items),
|
||||
"training_has_next_page": bool(isinstance(payload, dict) and payload.get("next")),
|
||||
"training_status_counts": dict(statuses),
|
||||
"training_models": models[:50],
|
||||
"training_samples": samples,
|
||||
}
|
||||
|
||||
|
||||
def probe_account_resources(key, proxy, timeout, debug=False):
|
||||
summaries = {}
|
||||
endpoint_statuses = {}
|
||||
model_refs = set()
|
||||
ok_count = 0
|
||||
total_items = 0
|
||||
summarizers = {
|
||||
"predictions": summarize_predictions,
|
||||
"deployments": summarize_deployments,
|
||||
"trainings": summarize_trainings,
|
||||
}
|
||||
for name, url in RESOURCE_ENDPOINTS.items():
|
||||
result = api_get(key, url, proxy, timeout, debug)
|
||||
endpoint_statuses[name] = {k: v for k, v in result.items() if k in ("status", "http_status", "message")}
|
||||
if result.get("status") != "OK":
|
||||
continue
|
||||
ok_count += 1
|
||||
summary = summarizers[name](result.get("payload"))
|
||||
summaries.update(summary)
|
||||
for key_name, value in summary.items():
|
||||
if key_name.endswith("_models") and isinstance(value, list):
|
||||
model_refs.update(value)
|
||||
total_items += sum(
|
||||
int(summary.get(field, 0) or 0)
|
||||
for field in ("prediction_count_sample", "deployment_count", "training_count_sample")
|
||||
)
|
||||
if ok_count == len(RESOURCE_ENDPOINTS):
|
||||
probe_status = "RESOURCE_OK" if total_items else "NO_RESOURCES"
|
||||
elif ok_count:
|
||||
probe_status = "PARTIAL"
|
||||
else:
|
||||
probe_status = next((item.get("status") for item in endpoint_statuses.values() if item.get("status")), "UNKNOWN")
|
||||
return {
|
||||
"probe": {"status": probe_status, "endpoints": endpoint_statuses},
|
||||
"models": sorted(model_refs)[:50],
|
||||
"model_count": len(model_refs),
|
||||
**summaries,
|
||||
}
|
||||
|
||||
|
||||
def iter_candidate_decisions(input_file, plain_files):
|
||||
seen_plain = set()
|
||||
for item in iter_findings(input_file, [DETECTOR]):
|
||||
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, args):
|
||||
account = api_get(key, ACCOUNT_URL, proxy, args.timeout, args.debug)
|
||||
if account.get("status") != "OK":
|
||||
return {k: v for k, v in account.items() if k != "payload"}
|
||||
data = account.get("payload") if isinstance(account.get("payload"), dict) else {}
|
||||
result = {
|
||||
"status": "VALID",
|
||||
"message": "account endpoint accepted",
|
||||
"account": data.get("username") or data.get("name") or "",
|
||||
"account_type": data.get("type") or "",
|
||||
}
|
||||
if not args.no_resource_probe:
|
||||
result.update(probe_account_resources(key, proxy, args.timeout, args.debug))
|
||||
result["message"] = (
|
||||
f"account endpoint accepted; probe={result.get('probe', {}).get('status')}; "
|
||||
f"models={result.get('model_count', 0)}; "
|
||||
f"deployments={result.get('deployment_count', 0)}; "
|
||||
f"predictions={result.get('prediction_count_sample', 0)}; "
|
||||
f"trainings={result.get('training_count_sample', 0)}"
|
||||
)
|
||||
else:
|
||||
result.update({"probe": {"status": "not_probed"}, "models": [], "model_count": 0})
|
||||
return result
|
||||
|
||||
|
||||
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 [])[:1000] if status == "VALID" else source
|
||||
commit_status_transaction(
|
||||
CHECKED_FILE, STATUS_FILES, key, status, result.get("message", ""), extra or source,
|
||||
)
|
||||
record_validation_result(SERVICE, key, result, source, finding, DETECTOR)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Replicate 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")
|
||||
parser.add_argument("--no-resource-probe", action="store_true", help="Only call /account; skip read-only predictions/deployments/trainings probes.")
|
||||
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 = 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("--- Replicate key checker ---")
|
||||
print("Default mode: /account plus read-only /predictions, /deployments and /trainings probes. Use --no-resource-probe for /account only.")
|
||||
print(f"proxy: {args.proxy_file}")
|
||||
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}] Replicate candidate {mask_secret(key)} from {source}")
|
||||
result = (
|
||||
check_key(key, next(proxy_cycler) if proxy_cycler else None, args)
|
||||
if valid_format else
|
||||
{"status": "NO_CONTEXT", "message": "candidate does not match canonical Replicate token format"}
|
||||
)
|
||||
print(f" STATUS: {result['status']} | {result.get('message', '')[:200]}")
|
||||
if result.get("status") == "VALID":
|
||||
print(f" ACCOUNT: {result.get('account') or 'unknown'}")
|
||||
print(f" MODELS: {result.get('model_count', 0)} from account resources")
|
||||
notable = result.get("models") or []
|
||||
if notable:
|
||||
print(f" MODEL REFS: {', '.join(notable[:8])}")
|
||||
print(f" PROBE: {(result.get('probe') or {}).get('status')}")
|
||||
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