Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,738 @@
|
||||
import sys
|
||||
|
||||
sys.dont_write_bytecode = True
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from itertools import cycle
|
||||
|
||||
import requests
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from keycheck_common import (
|
||||
acquire_file_lock,
|
||||
append_checked,
|
||||
default_input_file,
|
||||
default_proxy_file,
|
||||
env_int,
|
||||
ensure_output_files as ensure_private_output_files,
|
||||
iter_findings,
|
||||
iter_bounded_text_lines,
|
||||
keycheck_input_mode,
|
||||
load_known_statuses,
|
||||
private_atomic_writer,
|
||||
record_cached_keycheck_occurrence,
|
||||
record_validation_result,
|
||||
release_file_lock,
|
||||
require_provider_authority,
|
||||
service_output_dir,
|
||||
should_skip_key,
|
||||
write_keycheck_event,
|
||||
)
|
||||
from runtime_security import durable_replace, reject_reparse_components, require_private_directory, require_private_file
|
||||
|
||||
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
PARENT_DIR = os.path.dirname(SCRIPT_DIR)
|
||||
SERVICE = "gemini"
|
||||
OUTPUT_DIR = os.getenv("KEYCHECK_OUTPUT_DIR") or service_output_dir(SERVICE)
|
||||
|
||||
|
||||
def here(*parts):
|
||||
return os.path.join(SCRIPT_DIR, *parts)
|
||||
|
||||
|
||||
def out(*parts):
|
||||
return os.path.join(OUTPUT_DIR, *parts)
|
||||
|
||||
|
||||
def parent(*parts):
|
||||
return os.path.join(PARENT_DIR, *parts)
|
||||
|
||||
|
||||
# --- Configuration ---
|
||||
DEFAULT_INPUT_FILE = os.getenv("KEYCHECK_INPUT_FILE") or default_input_file()
|
||||
DEFAULT_PLAIN_INPUT_FILES = [out("gem.txt")]
|
||||
PROXY_FILE = os.getenv("KEYCHECK_PROXY_FILE") or default_proxy_file()
|
||||
|
||||
CHECKED_FILE = out("geminiChecked.txt")
|
||||
RESULTS_FILE = out("geminiResults.jsonl")
|
||||
|
||||
STATUS_FILES = {
|
||||
"VALID": out("geminiAlive.txt"),
|
||||
"VALID_RATE_LIMITED": out("geminiAliveRateLimited.txt"),
|
||||
"INVALID": out("geminiDead.txt"),
|
||||
"EXPIRED": out("geminiExpired.txt"),
|
||||
"LEAKED_REVOKED": out("geminiLeaked.txt"),
|
||||
"API_DISABLED": out("geminiDisabled.txt"),
|
||||
"RESTRICTED": out("geminiRestricted.txt"),
|
||||
"RATE_LIMITED": out("geminiRateLimited.txt"),
|
||||
"NETWORK_ERROR": out("geminiNetwork.txt"),
|
||||
"UNKNOWN": out("geminiUnknown.txt"),
|
||||
}
|
||||
|
||||
GEMINI_KEY_REGEX = re.compile(r"(?:AIza[0-9A-Za-z\-_]{35}|AQ\.[0-9A-Za-z\-_]{50})")
|
||||
GEMINI_DETECTOR_NAMES = {"googleai", "googleaistudio"}
|
||||
MODELS_URL = "https://generativelanguage.googleapis.com/v1beta/models"
|
||||
|
||||
PROBE_MODEL_PRIORITY = [
|
||||
"gemini-3.1-pro-preview",
|
||||
"gemini-3.7-flash",
|
||||
]
|
||||
|
||||
MODEL_PRIORITY = [
|
||||
"gemini-3",
|
||||
"gemini-2.5-pro",
|
||||
"gemini-2.5-flash",
|
||||
"gemini-2.0-flash",
|
||||
"gemini-1.5-pro",
|
||||
"gemini-1.5-flash",
|
||||
"imagen",
|
||||
"embedding",
|
||||
]
|
||||
|
||||
|
||||
def now_iso():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
def mask_key(key):
|
||||
if not key or len(key) < 12:
|
||||
return key
|
||||
return f"{key[:8]}...{key[-4:]}"
|
||||
|
||||
|
||||
def redact_key_text(text, key):
|
||||
if not isinstance(text, str):
|
||||
return text
|
||||
redacted = text.replace(key, "***REDACTED***") if key else text
|
||||
return GEMINI_KEY_REGEX.sub("***REDACTED***", redacted)
|
||||
|
||||
|
||||
def redact_result_text(result, key):
|
||||
if isinstance(result, dict):
|
||||
return {k: redact_result_text(v, key) for k, v in result.items()}
|
||||
if isinstance(result, list):
|
||||
return [redact_result_text(v, key) for v in result]
|
||||
return redact_key_text(result, key)
|
||||
|
||||
|
||||
def key_from_line(line):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
return None
|
||||
if "\t" in line:
|
||||
return line.split("\t", 1)[0].strip()
|
||||
return line.split(":", 1)[0].strip()
|
||||
|
||||
|
||||
def load_keys_from_file(filepath):
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return set()
|
||||
if not os.path.exists(filepath):
|
||||
return set()
|
||||
keys = set()
|
||||
for line in iter_bounded_text_lines(filepath):
|
||||
key = key_from_line(line)
|
||||
if key:
|
||||
keys.add(key)
|
||||
return keys
|
||||
|
||||
|
||||
def load_checked_statuses(filepath=CHECKED_FILE):
|
||||
statuses = {}
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return statuses
|
||||
if not os.path.exists(filepath):
|
||||
return statuses
|
||||
for line in iter_bounded_text_lines(filepath):
|
||||
parts = line.rstrip("\n").split("\t")
|
||||
if not parts or not parts[0]:
|
||||
continue
|
||||
key = parts[0]
|
||||
status = parts[1] if len(parts) > 1 else "UNKNOWN"
|
||||
statuses[key] = status
|
||||
return statuses
|
||||
|
||||
|
||||
def load_all_known_keys():
|
||||
known = set(load_checked_statuses().keys())
|
||||
for path in STATUS_FILES.values():
|
||||
known.update(load_keys_from_file(path))
|
||||
return known
|
||||
|
||||
|
||||
def ensure_output_files():
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return
|
||||
require_private_directory(OUTPUT_DIR, create=True)
|
||||
|
||||
legacy_rate_limited = out("geminiLimited.txt")
|
||||
rate_limited = STATUS_FILES["RATE_LIMITED"]
|
||||
if os.path.exists(legacy_rate_limited) and not os.path.exists(rate_limited):
|
||||
require_private_file(legacy_rate_limited)
|
||||
durable_replace(legacy_rate_limited, rate_limited)
|
||||
require_private_file(rate_limited)
|
||||
|
||||
paths = {CHECKED_FILE, RESULTS_FILE, *STATUS_FILES.values()}
|
||||
ensure_private_output_files(paths)
|
||||
migrate_legacy_alive_rate_limited()
|
||||
|
||||
|
||||
def effective_status(result):
|
||||
status = result.get("status")
|
||||
probe_status = (result.get("probe") or {}).get("status")
|
||||
if status == "VALID" and probe_status == "RATE_LIMITED":
|
||||
return "VALID_RATE_LIMITED"
|
||||
return status
|
||||
|
||||
|
||||
def _gemini_status_layout():
|
||||
paths_by_status = {
|
||||
status: os.path.abspath(os.fspath(path))
|
||||
for status, path in STATUS_FILES.items()
|
||||
}
|
||||
paths = list(dict.fromkeys(paths_by_status.values()))
|
||||
directories = {os.path.normcase(os.path.dirname(path)) for path in paths}
|
||||
if len(paths) != len(paths_by_status) or len(directories) != 1:
|
||||
raise RuntimeError("Gemini status files must be unique files in one directory")
|
||||
directory = os.path.dirname(paths[0])
|
||||
require_private_directory(directory, create=True)
|
||||
return paths_by_status, paths, os.path.join(directory, "geminiStatus.lock")
|
||||
|
||||
|
||||
def migrate_legacy_alive_rate_limited():
|
||||
paths_by_status, _, lock_path = _gemini_status_layout()
|
||||
alive_path = paths_by_status["VALID"]
|
||||
limited_path = paths_by_status["VALID_RATE_LIMITED"]
|
||||
lock = acquire_file_lock(lock_path, timeout_sec=30)
|
||||
try:
|
||||
if not os.path.lexists(alive_path):
|
||||
return
|
||||
require_private_file(alive_path)
|
||||
|
||||
keep = []
|
||||
moved = {}
|
||||
for line in iter_bounded_text_lines(alive_path):
|
||||
key = key_from_line(line)
|
||||
if key and line.rstrip("\r\n").endswith(":RATE_LIMITED"):
|
||||
moved.setdefault(key, line if line.endswith("\n") else f"{line}\n")
|
||||
else:
|
||||
keep.append(line)
|
||||
if not moved:
|
||||
return
|
||||
|
||||
if os.path.lexists(limited_path):
|
||||
require_private_file(limited_path)
|
||||
existing = _normalized_status_lines(list(iter_bounded_text_lines(limited_path)))
|
||||
limited = []
|
||||
published = set()
|
||||
for line in existing:
|
||||
key = key_from_line(line)
|
||||
if key in moved:
|
||||
if key in published:
|
||||
continue
|
||||
published.add(key)
|
||||
limited.append(line)
|
||||
for key, line in moved.items():
|
||||
if key not in published:
|
||||
limited.append(line)
|
||||
published.add(key)
|
||||
|
||||
_validate_status_snapshot(limited_path, limited)
|
||||
_validate_status_snapshot(alive_path, keep)
|
||||
|
||||
# Make every moved key durable before publishing the source snapshot
|
||||
# that removes it. An interruption can therefore only leave duplicates.
|
||||
_replace_status_snapshot(limited_path, limited)
|
||||
confirmed = {key: 0 for key in moved}
|
||||
for line in iter_bounded_text_lines(limited_path):
|
||||
key = key_from_line(line)
|
||||
if key in confirmed:
|
||||
confirmed[key] += 1
|
||||
if any(count != 1 for count in confirmed.values()):
|
||||
raise RuntimeError("Gemini legacy rate-limited status publication was incomplete")
|
||||
_replace_status_snapshot(alive_path, keep)
|
||||
finally:
|
||||
release_file_lock(lock, lock_path)
|
||||
|
||||
|
||||
def load_proxies(proxy_file):
|
||||
if not os.path.exists(proxy_file):
|
||||
print(f"Info: {proxy_file} not found. Requests will go directly.")
|
||||
return None
|
||||
|
||||
proxies = []
|
||||
with open(proxy_file, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
ip, port, login, password = line.split(":")
|
||||
proxy_url = f"http://{login}:{password}@{ip}:{port}"
|
||||
proxies.append({"http": proxy_url, "https": proxy_url})
|
||||
except ValueError:
|
||||
print(f"Warning: bad proxy format: {line}. Skipping.")
|
||||
|
||||
if not proxies:
|
||||
print(f"Warning: {proxy_file} is empty. Requests will go directly.")
|
||||
return None
|
||||
|
||||
print(f"Loaded proxies: {len(proxies)}")
|
||||
return cycle(proxies)
|
||||
|
||||
|
||||
def status_file_line(key, result, status):
|
||||
if status in ("VALID", "VALID_RATE_LIMITED"):
|
||||
models_str = ",".join(result.get("notable_models", [])) or "models-only"
|
||||
probe_status = result.get("probe", {}).get("status", "not_probed")
|
||||
return f"{key}:[{models_str}]:{result.get('model_class', 'unknown')}:{probe_status}\n"
|
||||
message = (result.get("error", {}).get("message") or "").replace("\n", " ")[:300]
|
||||
return f"{key}\t{status}\t{message}\n"
|
||||
|
||||
|
||||
def _normalized_status_lines(lines):
|
||||
return [line if line.endswith("\n") else f"{line}\n" for line in lines]
|
||||
|
||||
|
||||
def _validate_status_snapshot(path, lines):
|
||||
max_bytes = max(1, env_int("KEYCHECK_INPUT_LIST_MAX_BYTES", 32 * 1024 * 1024))
|
||||
max_items = max(1, env_int("KEYCHECK_INPUT_LIST_MAX_ITEMS", 100000))
|
||||
max_line_bytes = max(1, env_int("KEYCHECK_INPUT_LIST_MAX_LINE_BYTES", 8192))
|
||||
if len(lines) > max_items:
|
||||
raise RuntimeError(f"Gemini status file exceeds its item bound: {path}")
|
||||
total = 0
|
||||
for index, line in enumerate(lines, 1):
|
||||
encoded = line.encode("utf-8")
|
||||
if len(encoded) > max_line_bytes:
|
||||
raise RuntimeError(f"Gemini status line exceeds its byte bound: {path}:{index}")
|
||||
total += len(encoded)
|
||||
if total > max_bytes:
|
||||
raise RuntimeError(f"Gemini status file exceeds its aggregate byte bound: {path}")
|
||||
|
||||
|
||||
def _replace_status_snapshot(path, lines):
|
||||
with private_atomic_writer(path, binary=True, suffix=".status.tmp") as handle:
|
||||
for line in lines:
|
||||
handle.write(line.encode("utf-8"))
|
||||
|
||||
|
||||
def append_status_file(key, result):
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return
|
||||
paths_by_status, paths, lock_path = _gemini_status_layout()
|
||||
lock = acquire_file_lock(lock_path, timeout_sec=30)
|
||||
try:
|
||||
status = effective_status(result)
|
||||
target_path = paths_by_status.get(status, paths_by_status["UNKNOWN"])
|
||||
new_line = status_file_line(key, result, status)
|
||||
snapshots = {}
|
||||
for path in paths:
|
||||
if os.path.lexists(path):
|
||||
reject_reparse_components(path)
|
||||
snapshots[path] = _normalized_status_lines(list(iter_bounded_text_lines(path)))
|
||||
|
||||
rewritten = {
|
||||
path: [line for line in lines if key_from_line(line) != key]
|
||||
for path, lines in snapshots.items()
|
||||
}
|
||||
rewritten[target_path].insert(0, new_line)
|
||||
for path, lines in rewritten.items():
|
||||
_validate_status_snapshot(path, lines)
|
||||
|
||||
# Publish the new classification before removing any old copies. A
|
||||
# failure after this point can leave duplicates, but never no status.
|
||||
_replace_status_snapshot(target_path, rewritten[target_path])
|
||||
for path in paths:
|
||||
if path == target_path or rewritten[path] == snapshots[path]:
|
||||
continue
|
||||
_replace_status_snapshot(path, rewritten[path])
|
||||
finally:
|
||||
release_file_lock(lock, lock_path)
|
||||
|
||||
|
||||
def append_checked_file(key, result):
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return
|
||||
append_checked(CHECKED_FILE, key, effective_status(result))
|
||||
|
||||
|
||||
def is_gemini_detector(detector):
|
||||
return str(detector or "").lower() in GEMINI_DETECTOR_NAMES
|
||||
|
||||
|
||||
def custom_detector_name(data):
|
||||
if not isinstance(data, dict):
|
||||
return ""
|
||||
extra = data.get("ExtraData") if isinstance(data.get("ExtraData"), dict) else {}
|
||||
name = extra.get("name") or ""
|
||||
if str(data.get("DetectorName") or "").lower() == "customregex" and is_gemini_detector(name):
|
||||
return name
|
||||
return ""
|
||||
|
||||
|
||||
def detector_name_from_finding(data):
|
||||
if not isinstance(data, dict):
|
||||
return ""
|
||||
if is_gemini_detector(data.get("DetectorName")):
|
||||
return data.get("DetectorName")
|
||||
custom_name = custom_detector_name(data)
|
||||
if custom_name:
|
||||
return custom_name
|
||||
|
||||
# Old wrapped format from earlier scanner versions.
|
||||
if is_gemini_detector(data.get("detector")):
|
||||
return data.get("detector")
|
||||
|
||||
finding = data.get("finding")
|
||||
if isinstance(finding, dict) and is_gemini_detector(finding.get("DetectorName")):
|
||||
return finding.get("DetectorName")
|
||||
custom_name = custom_detector_name(finding)
|
||||
if custom_name:
|
||||
return custom_name
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def extract_key_from_finding(data):
|
||||
if is_gemini_detector(data.get("DetectorName")):
|
||||
return data.get("Raw") or data.get("RawV2")
|
||||
if custom_detector_name(data):
|
||||
return data.get("Raw") or data.get("RawV2")
|
||||
|
||||
# Old wrapped format from earlier scanner versions.
|
||||
if is_gemini_detector(data.get("detector")):
|
||||
return data.get("raw") or data.get("raw_v2")
|
||||
|
||||
finding = data.get("finding")
|
||||
if isinstance(finding, dict) and is_gemini_detector(finding.get("DetectorName")):
|
||||
return finding.get("Raw") or finding.get("RawV2")
|
||||
if custom_detector_name(finding):
|
||||
return finding.get("Raw") or finding.get("RawV2")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def iter_candidate_keys(input_file, plain_files):
|
||||
for item in iter_findings(input_file, ["GoogleAI", "GoogleAIStudio", "CustomRegex"]):
|
||||
finding = item.get("finding") or {}
|
||||
key = item.get("raw") or extract_key_from_finding(finding)
|
||||
if key and GEMINI_KEY_REGEX.fullmatch(key) and detector_name_from_finding(finding):
|
||||
yield item.get("source") or input_file, key, finding
|
||||
|
||||
if keycheck_input_mode() == 'postgres':
|
||||
return
|
||||
for path in plain_files:
|
||||
if not os.path.exists(path):
|
||||
print(f"Info: plain input {path} not found. Skipping.")
|
||||
continue
|
||||
try:
|
||||
keys = set()
|
||||
for line in iter_bounded_text_lines(path):
|
||||
keys.update(GEMINI_KEY_REGEX.findall(line))
|
||||
except (OSError, RuntimeError) as e:
|
||||
print(f"Warning: cannot read {path}: {e}")
|
||||
continue
|
||||
for idx, key in enumerate(sorted(keys), 1):
|
||||
yield f"{path}:plain:{idx}", key, {}
|
||||
|
||||
|
||||
def parse_error_response(response):
|
||||
try:
|
||||
payload = response.json()
|
||||
except json.JSONDecodeError:
|
||||
payload = {}
|
||||
|
||||
error = payload.get("error", {}) if isinstance(payload, dict) else {}
|
||||
return {
|
||||
"http_status": response.status_code,
|
||||
"code": error.get("code", response.status_code),
|
||||
"status": error.get("status", ""),
|
||||
"message": error.get("message", response.text[:500]),
|
||||
}
|
||||
|
||||
|
||||
def classify_error(error):
|
||||
http_status = int(error.get("http_status") or 0)
|
||||
status = str(error.get("status") or "").lower()
|
||||
message = str(error.get("message") or "").lower()
|
||||
|
||||
if "reported as leaked" in message or "leaked" in message:
|
||||
return "LEAKED_REVOKED"
|
||||
if "api key expired" in message or "expired" in message:
|
||||
return "EXPIRED"
|
||||
if "api key not valid" in message or "invalid api key" in message:
|
||||
return "INVALID"
|
||||
if "has not been used" in message or "it is disabled" in message or "api is disabled" in message:
|
||||
return "API_DISABLED"
|
||||
if "requests to this api" in message and "blocked" in message:
|
||||
return "RESTRICTED"
|
||||
if "api key restrictions" in message or "permission_denied" in status:
|
||||
return "RESTRICTED"
|
||||
if http_status == 429 or "resource_exhausted" in status or "quota" in message:
|
||||
return "RATE_LIMITED"
|
||||
if http_status in (400, 401):
|
||||
return "INVALID"
|
||||
if http_status == 403:
|
||||
return "RESTRICTED"
|
||||
return "UNKNOWN"
|
||||
|
||||
|
||||
def fetch_models(key, proxy, timeout, debug=False):
|
||||
try:
|
||||
response = requests.get(MODELS_URL, params={"key": key}, proxies=proxy, timeout=timeout)
|
||||
except requests.RequestException as e:
|
||||
return {
|
||||
"status": "NETWORK_ERROR",
|
||||
"error": {"message": str(e)},
|
||||
"models": [],
|
||||
"model_infos": [],
|
||||
}
|
||||
|
||||
if debug:
|
||||
print(f" DEBUG /models: HTTP {response.status_code}: {redact_key_text(response.text[:500], key)}")
|
||||
|
||||
if response.status_code != 200:
|
||||
error = parse_error_response(response)
|
||||
return {
|
||||
"status": classify_error(error),
|
||||
"error": error,
|
||||
"models": [],
|
||||
"model_infos": [],
|
||||
}
|
||||
|
||||
payload = response.json()
|
||||
model_infos = payload.get("models", [])
|
||||
models = sorted({m.get("name", "").replace("models/", "") for m in model_infos if m.get("name")})
|
||||
return {
|
||||
"status": "VALID",
|
||||
"error": {},
|
||||
"models": models,
|
||||
"model_infos": model_infos,
|
||||
}
|
||||
|
||||
|
||||
def supported_methods_by_model(model_infos):
|
||||
output = {}
|
||||
for model in model_infos:
|
||||
name = model.get("name", "").replace("models/", "")
|
||||
if not name:
|
||||
continue
|
||||
output[name] = sorted(model.get("supportedGenerationMethods", []))
|
||||
return output
|
||||
|
||||
|
||||
def classify_models(models, methods_by_model):
|
||||
notable = []
|
||||
lower_models = {m.lower(): m for m in models}
|
||||
for marker in MODEL_PRIORITY:
|
||||
for lower, original in lower_models.items():
|
||||
if marker in lower and original not in notable:
|
||||
notable.append(original)
|
||||
|
||||
generation_models = sorted([
|
||||
model for model, methods in methods_by_model.items()
|
||||
if "generateContent" in methods
|
||||
])
|
||||
|
||||
if any("gemini-2.5-pro" in m.lower() for m in generation_models):
|
||||
model_class = "pro_generation"
|
||||
elif generation_models:
|
||||
model_class = "generation"
|
||||
elif models:
|
||||
model_class = "models_only"
|
||||
else:
|
||||
model_class = "no_models"
|
||||
|
||||
return notable[:20], generation_models, model_class
|
||||
|
||||
|
||||
def choose_probe_model(generation_models):
|
||||
available = set(generation_models)
|
||||
for model in PROBE_MODEL_PRIORITY:
|
||||
if model in available:
|
||||
return model
|
||||
return generation_models[0] if generation_models else None
|
||||
|
||||
|
||||
def probe_generation(key, model, proxy, timeout, debug=False):
|
||||
if not model:
|
||||
return {"status": "NO_GENERATION_MODEL", "model": None}
|
||||
|
||||
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
|
||||
headers = {"x-goog-api-key": key, "Content-Type": "application/json"}
|
||||
payload = {
|
||||
"contents": [{"parts": [{"text": "ping"}]}],
|
||||
"generationConfig": {"maxOutputTokens": 1},
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(url, headers=headers, json=payload, proxies=proxy, timeout=timeout)
|
||||
except requests.RequestException as e:
|
||||
return {"status": "NETWORK_ERROR", "model": model, "error": {"message": str(e)}}
|
||||
|
||||
if debug:
|
||||
print(f" DEBUG probe {model}: HTTP {response.status_code}: {redact_key_text(response.text[:500], key)}")
|
||||
|
||||
if response.status_code == 200:
|
||||
return {"status": "GENERATION_OK", "model": model}
|
||||
|
||||
error = parse_error_response(response)
|
||||
return {"status": classify_error(error), "model": model, "error": error}
|
||||
|
||||
|
||||
def check_key(key, proxy, args):
|
||||
result = fetch_models(key, proxy, args.timeout, args.debug)
|
||||
result = redact_result_text(result, key)
|
||||
result.update({
|
||||
"checked_at": now_iso(),
|
||||
"key_masked": mask_key(key),
|
||||
"model_count": len(result.get("models", [])),
|
||||
})
|
||||
|
||||
if result["status"] != "VALID":
|
||||
result["notable_models"] = []
|
||||
result["generation_models"] = []
|
||||
result["model_class"] = "none"
|
||||
return result
|
||||
|
||||
methods_by_model = supported_methods_by_model(result.get("model_infos", []))
|
||||
notable, generation_models, model_class = classify_models(result["models"], methods_by_model)
|
||||
|
||||
result["methods_by_model"] = methods_by_model
|
||||
result["notable_models"] = notable
|
||||
result["generation_models"] = generation_models[:50]
|
||||
result["model_class"] = model_class
|
||||
result["billing_status"] = "unknown"
|
||||
|
||||
if args.probe_generation:
|
||||
probe_model = choose_probe_model(generation_models)
|
||||
result["probe"] = redact_result_text(probe_generation(key, probe_model, proxy, args.timeout, args.debug), key)
|
||||
else:
|
||||
result["probe"] = {"status": "not_probed", "model": None}
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def print_result(index, source, key, result):
|
||||
print(f"\n[{index}] Candidate {mask_key(key)} from {source}")
|
||||
print(f" STATUS: {result['status']}")
|
||||
|
||||
if result["status"] == "VALID":
|
||||
print(f" MODELS: {result.get('model_count', 0)} total; class={result.get('model_class')}")
|
||||
notable = result.get("notable_models", [])[:8]
|
||||
if notable:
|
||||
print(f" NOTABLE: {', '.join(notable)}")
|
||||
probe = result.get("probe", {})
|
||||
print(f" PROBE: {probe.get('status')} ({probe.get('model')})")
|
||||
if effective_status(result) == "VALID_RATE_LIMITED":
|
||||
print(f" OUT: {STATUS_FILES['VALID_RATE_LIMITED']}")
|
||||
else:
|
||||
error = result.get("error", {})
|
||||
message = (error.get("message") or "").replace("\n", " ")[:300]
|
||||
if message:
|
||||
print(f" MESSAGE: {message}")
|
||||
|
||||
print(f" OUT: {STATUS_FILES.get(result['status'], STATUS_FILES['UNKNOWN'])}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gemini / Google AI API key classifier")
|
||||
parser.add_argument("--input", default=DEFAULT_INPUT_FILE)
|
||||
parser.add_argument("--plain", action="append", default=None, help="Plain text file with Gemini keys. Can be repeated.")
|
||||
parser.add_argument("--proxy-file", default=PROXY_FILE)
|
||||
parser.add_argument("--timeout", type=int, default=20)
|
||||
parser.add_argument("--max-keys", type=int, default=0)
|
||||
parser.add_argument("--probe-generation", action="store_true", help="Optionally call generateContent, preferring gemini-3.1-pro-preview when available.")
|
||||
parser.add_argument("--retry-limited", action="store_true")
|
||||
parser.add_argument("--retry-unknown", action="store_true")
|
||||
parser.add_argument("--retry-network", 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()
|
||||
plain_files = args.plain if args.plain is not None else DEFAULT_PLAIN_INPUT_FILES
|
||||
ensure_output_files()
|
||||
|
||||
print("--- Gemini key checker ---")
|
||||
print("Default mode: /models only. Use --probe-generation for runtime/billing probe.")
|
||||
print(f"Workspace: {SCRIPT_DIR}")
|
||||
|
||||
proxy_cycler = load_proxies(args.proxy_file)
|
||||
checked_statuses = load_checked_statuses()
|
||||
known_statuses = load_known_statuses(CHECKED_FILE, STATUS_FILES)
|
||||
known_keys = set(known_statuses)
|
||||
retry_statuses = set()
|
||||
if args.retry_limited:
|
||||
retry_statuses.update(("RATE_LIMITED", "VALID_RATE_LIMITED"))
|
||||
if args.retry_unknown:
|
||||
retry_statuses.add("UNKNOWN")
|
||||
if args.retry_network:
|
||||
retry_statuses.add("NETWORK_ERROR")
|
||||
if args.retry_valid:
|
||||
retry_statuses.update(("VALID", "VALID_RATE_LIMITED"))
|
||||
print(f"Loaded known keys: {len(known_keys)}; checked records: {len(checked_statuses)}")
|
||||
|
||||
seen_this_run = set()
|
||||
processed = 0
|
||||
skipped = 0
|
||||
|
||||
for source, key, finding in iter_candidate_keys(args.input, plain_files):
|
||||
if keycheck_input_mode() != 'postgres' and key in seen_this_run:
|
||||
cached_status = checked_statuses.get(key) or known_statuses.get(key) or 'UNKNOWN'
|
||||
detector = detector_name_from_finding(finding) or "GoogleAI"
|
||||
record_cached_keycheck_occurrence(SERVICE, key, cached_status, source, finding, detector)
|
||||
skipped += 1
|
||||
continue
|
||||
seen_this_run.add(key)
|
||||
|
||||
detector = detector_name_from_finding(finding) or "GoogleAI"
|
||||
if should_skip_key(
|
||||
key, checked_statuses, known_keys, args, retry_statuses,
|
||||
service=SERVICE, source=source, finding=finding, detector=detector,
|
||||
known_statuses=known_statuses,
|
||||
):
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
if args.max_keys and processed >= args.max_keys:
|
||||
break
|
||||
|
||||
processed += 1
|
||||
|
||||
proxy = next(proxy_cycler) if proxy_cycler else None
|
||||
result = check_key(key, proxy, args)
|
||||
result["source"] = source
|
||||
event_result = {**result, "status": effective_status(result)}
|
||||
write_keycheck_event(SERVICE, RESULTS_FILE, key, event_result, source, finding, detector)
|
||||
|
||||
print_result(processed, source, key, result)
|
||||
append_status_file(key, result)
|
||||
append_checked_file(key, result)
|
||||
record_validation_result(SERVICE, key, {**result, "status": effective_status(result)}, source, finding, detector)
|
||||
|
||||
known_keys.add(key)
|
||||
checked_statuses[key] = effective_status(result)
|
||||
|
||||
# Small pause helps when many keys hit the same API/proxy.
|
||||
time.sleep(0.1)
|
||||
|
||||
print("\n--- Done ---")
|
||||
print(f"Processed: {processed}")
|
||||
print(f"Skipped: {skipped}")
|
||||
print(f"Results: {RESULTS_FILE}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user