Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,203 @@
|
||||
import sys
|
||||
|
||||
sys.dont_write_bytecode = True
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from keycheck_common import (
|
||||
commit_status_transaction,
|
||||
default_input_file,
|
||||
default_proxy_file,
|
||||
ensure_output_files,
|
||||
iter_findings,
|
||||
load_checked_statuses,
|
||||
load_known_keys,
|
||||
load_proxies,
|
||||
mask_secret,
|
||||
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 = "provider_resolver"
|
||||
DETECTOR = "ProviderResolver"
|
||||
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, "providerResolverChecked.txt")
|
||||
RESULTS_FILE = os.path.join(OUTPUT_DIR, "providerResolverResults.jsonl")
|
||||
STATUS_FILES = {
|
||||
"VALID": os.path.join(OUTPUT_DIR, "providerResolverAlive.txt"),
|
||||
"NO_BALANCE": os.path.join(OUTPUT_DIR, "providerResolverNoBalance.txt"),
|
||||
"DEAD": os.path.join(OUTPUT_DIR, "providerResolverDead.txt"),
|
||||
"RESTRICTED": os.path.join(OUTPUT_DIR, "providerResolverRestricted.txt"),
|
||||
"LIMITED": os.path.join(OUTPUT_DIR, "providerResolverLimited.txt"),
|
||||
"NETWORK": os.path.join(OUTPUT_DIR, "providerResolverNetwork.txt"),
|
||||
"NO_CONTEXT": os.path.join(OUTPUT_DIR, "providerResolverNoContext.txt"),
|
||||
"UNKNOWN": os.path.join(OUTPUT_DIR, "providerResolverUnknown.txt"),
|
||||
}
|
||||
|
||||
RESOLVABLE_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_.-])"
|
||||
)
|
||||
FOREIGN_KEY_PREFIXES = ("sk-ant-", "sk-or-", "sk-proj-", "sk-svcacct-", "sk-admin-")
|
||||
AMBIGUOUS_HINTS = {AMBIGUOUS_QWEN_DEEPSEEK_HINT, AMBIGUOUS_GENERIC_SK_HINT}
|
||||
|
||||
|
||||
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 RESOLVABLE_KEY_REGEX.fullmatch(value):
|
||||
return "candidate does not match a bounded resolvable provider-key format"
|
||||
return ""
|
||||
|
||||
|
||||
def iter_candidate_keys(input_file):
|
||||
detector_names = [
|
||||
"ProviderResolver", "CustomRegex", "QwenDashScope", "Qwen_DashScope",
|
||||
"Qwen", "DashScope", "DeepSeek", "DeepSeekApiKey", "DeepSeek_API_Key",
|
||||
"KimiMoonshot", "MoonshotAI", "Moonshot", "Kimi", "ZaiGLM",
|
||||
"qwendashscope", "qwen_dashscope", "qwen", "dashscope", "deepseek",
|
||||
"deepseekapikey", "deepseek_api_key", "kimimoonshot", "moonshotai",
|
||||
"moonshot", "kimi", "zaiglm",
|
||||
]
|
||||
for item in iter_findings(input_file, detector_names):
|
||||
finding = item.get("finding") or {}
|
||||
key = item.get("credential_secret_text") or ""
|
||||
if not key:
|
||||
for value in (item.get("raw"), item.get("raw_v2"), finding.get("Raw"), finding.get("RawV2")):
|
||||
match = RESOLVABLE_KEY_REGEX.search(str(value or ""))
|
||||
if match:
|
||||
key = match.group(0)
|
||||
break
|
||||
if not key or key_rejection_reason(key):
|
||||
continue
|
||||
context = finding.get("ScannerContext") if isinstance(finding.get("ScannerContext"), dict) else {}
|
||||
metadata_hint = ""
|
||||
active_metadata = item.get("candidate_metadata")
|
||||
if isinstance(active_metadata, dict):
|
||||
metadata_hint = str(active_metadata.get("provider_hint") or "")
|
||||
metadata_candidates = active_metadata.get("provider_candidates")
|
||||
if metadata_hint or isinstance(metadata_candidates, list):
|
||||
context = dict(context)
|
||||
if metadata_hint:
|
||||
context.setdefault("provider_hint", metadata_hint)
|
||||
if isinstance(metadata_candidates, list):
|
||||
context.setdefault("provider_candidates", metadata_candidates)
|
||||
finding = dict(finding)
|
||||
finding["ScannerContext"] = context
|
||||
hint = str(context.get("provider_hint") or metadata_hint or AMBIGUOUS_GENERIC_SK_HINT)
|
||||
if hint not in AMBIGUOUS_HINTS:
|
||||
hint = AMBIGUOUS_GENERIC_SK_HINT
|
||||
yield key, item.get("source") or input_file, finding, hint
|
||||
|
||||
|
||||
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("resolved_provider") 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 parse_args():
|
||||
parser = argparse.ArgumentParser(description="Ambiguous generic provider key resolver")
|
||||
parser.add_argument("--input", default=INPUT_FILE)
|
||||
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("--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)
|
||||
processed = 0
|
||||
skipped = 0
|
||||
for key, source, finding, hint in iter_candidate_keys(args.input):
|
||||
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}] Ambiguous provider candidate {mask_secret(key)} from {source}")
|
||||
proxy = next(proxy_cycler) if proxy_cycler else None
|
||||
result = resolve_provider_key(
|
||||
key, finding, proxy, args.timeout, args.debug, hint=hint,
|
||||
)
|
||||
print(
|
||||
f" STATUS: {result['status']} provider={result.get('resolved_provider', '')} "
|
||||
f"| {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