sashatrask
143 lines
6.1 KiB
Python
143 lines
6.1 KiB
Python
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 = "huggingface"
|
|
DETECTOR_NAMES = ["HuggingFace", "Huggingface"]
|
|
DETECTOR = "HuggingFace"
|
|
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, "huggingfaceChecked.txt")
|
|
RESULTS_FILE = os.path.join(OUTPUT_DIR, "huggingfaceResults.jsonl")
|
|
STATUS_FILES = {
|
|
"VALID": os.path.join(OUTPUT_DIR, "huggingfaceAlive.txt"),
|
|
"DEAD": os.path.join(OUTPUT_DIR, "huggingfaceDead.txt"),
|
|
"RESTRICTED": os.path.join(OUTPUT_DIR, "huggingfaceRestricted.txt"),
|
|
"LIMITED": os.path.join(OUTPUT_DIR, "huggingfaceLimited.txt"),
|
|
"NETWORK": os.path.join(OUTPUT_DIR, "huggingfaceNetwork.txt"),
|
|
"NO_CONTEXT": os.path.join(OUTPUT_DIR, "huggingfaceNoContext.txt"),
|
|
"UNKNOWN": os.path.join(OUTPUT_DIR, "huggingfaceUnknown.txt"),
|
|
}
|
|
KEY_REGEX = re.compile(r"\bhf_[A-Za-z0-9]{20,}\b")
|
|
WHOAMI_URL = "https://huggingface.co/api/whoami-v2"
|
|
|
|
|
|
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(WHOAMI_URL, headers={"Authorization": f"Bearer {key}"}, 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 {}
|
|
return {"status": "VALID", "message": "whoami accepted", "username": data.get("name") or data.get("fullname") or ""}
|
|
status = "RESTRICTED" if response.status_code == 403 else classify_common_http_status(response.status_code)
|
|
return {"status": status, "http_status": response.status_code, "message": request_error_message(response).replace(key, "***REDACTED***")}
|
|
|
|
|
|
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", ""), source,
|
|
)
|
|
record_validation_result(SERVICE, key, result, source, finding, DETECTOR)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="HuggingFace 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("--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")
|
|
processed = skipped = 0
|
|
print("--- HuggingFace 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}] HuggingFace 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 Hugging Face 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()
|