#!/usr/bin/env python3
"""UNIVERSAL KEY CLASSIFIER — probes unknown sk-*/key-like tokens against
every known provider's cheapest endpoint to identify what they actually are.

Problem it solves: DeepSeek / DashScope / Kimi / many relays share the sk-<hex>
shape. A regex cannot tell them apart — only a live probe can. This module
takes a candidate key and returns which provider(s) it is ALIVE on + models.

Used as a second validation pass for keys whose DB provider came back DEAD —
often the key was misclassified and is alive on another endpoint.

Run standalone:  python universal_validator.py keys.txt [relay_base_url]
"""
import concurrent.futures as cf
import json
import sys
import urllib.error
import urllib.parse
import urllib.request

UA = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'}

# Direct-only opener: no system/registry proxies (urllib honors IE/registry
# proxy settings; lanes' S.trust_env=False is the requests-world equivalent).
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))

# (provider, base_url, header_fn, probe_path)
ENDPOINTS = [
    ('deepseek',   'https://api.deepseek.com',                          lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('dashscope',  'https://dashscope.aliyuncs.com/compatible-mode/v1', lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('moonshot',   'https://api.moonshot.cn/v1',                        lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('openai',     'https://api.openai.com/v1',                         lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('openrouter', 'https://openrouter.ai/api/v1',                      lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('groq',       'https://api.groq.com/openai/v1',                    lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('zai',        'https://api.z.ai/api/paas/v4',                      lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('mistral',    'https://api.mistral.ai/v1',                         lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('together',   'https://api.together.xyz/v1',                       lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('fireworks',  'https://api.fireworks.ai/v1',                       lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('anthropic',  'https://api.anthropic.com/v1',                      lambda k: {'x-api-key': k, 'anthropic-version': '2023-06-01'}, '/models'),
    ('siliconflow','https://api.siliconflow.cn/v1',                     lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
    ('novita',     'https://api.novita.ai/v3/openai',                   lambda k: {'Authorization': f'Bearer {k}'}, '/models'),
]


def _get(url, headers, timeout=12):
    h = dict(UA)
    h.update(headers)
    req = urllib.request.Request(url, headers=h)
    try:
        with _OPENER.open(req, timeout=timeout) as r:
            return r.status, r.read()
    except urllib.error.HTTPError as e:
        return e.code, e.read()[:400]
    except Exception:
        return -1, b''


def _chat_ok(base, headers, timeout=12, model='qwen-plus'):
    """1-token chat completion — the only real proof a key can serve inference."""
    payload = json.dumps({'model': model, 'messages': [{'role': 'user', 'content': 'hi'}], 'max_tokens': 1}).encode()
    h = dict(UA)
    h.update(headers)
    h['Content-Type'] = 'application/json'
    req = urllib.request.Request(base + '/chat/completions', data=payload, headers=h, method='POST')
    try:
        with _OPENER.open(req, timeout=timeout) as r:
            return r.status == 200
    except Exception:
        return False


# Auth-gated probe paths. Some providers expose /models publicly (OpenRouter,
# Novita) — probing those returns 200 even for a garbage key. So we probe an
# endpoint that REQUIRES a valid key, and treat 200 as alive.
AUTH_PROBE = {
    # provider: (base, path, header_fn, parser)
    'deepseek':   ('https://api.deepseek.com', '/user/balance', lambda k: {'Authorization': f'Bearer {k}'}),
    'dashscope':  ('https://dashscope.aliyuncs.com/compatible-mode/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'moonshot':   ('https://api.moonshot.cn/v1', '/users/me/balance', lambda k: {'Authorization': f'Bearer {k}'}),
    'openai':     ('https://api.openai.com/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'openrouter': ('https://openrouter.ai/api/v1', '/key', lambda k: {'Authorization': f'Bearer {k}'}),
    'groq':       ('https://api.groq.com/openai/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'zai':        ('https://api.z.ai/api/paas/v4', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'mistral':    ('https://api.mistral.ai/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'together':   ('https://api.together.xyz/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'fireworks':  ('https://api.fireworks.ai/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    'anthropic':  ('https://api.anthropic.com/v1', '/models', lambda k: {'x-api-key': k, 'anthropic-version': '2023-06-01'}),
    'siliconflow':('https://api.siliconflow.cn/v1', '/models', lambda k: {'Authorization': f'Bearer {k}'}),
    # novita removed: /models is public (200 for any key) — no auth-gated free endpoint
}

# Providers whose /models is PUBLIC — must use an auth-gated path instead.
PUBLIC_MODELS = {'openrouter', 'novita'}  # never trust their /models


def _parse_models(d):
    data = d.get('data') or d.get('models') or []
    return [m.get('id') or m.get('name', '') for m in data[:12]] if isinstance(data, list) else []


def probe(key, timeout=12):
    """Return providers where the key is ALIVE, using auth-gated endpoints only."""
    alive = []
    for name, (base, path, hdr_fn) in AUTH_PROBE.items():
        if name == 'dashscope':
            # /models returns 200 even for arrears (dead-balance) keys —
            # only a chat completion proves the key can serve inference.
            if _chat_ok(base, hdr_fn(key), timeout):
                alive.append({'provider': name, 'models': ['qwen — 260+ моделей DashScope']})
            continue
        st, body = _get(base + path, hdr_fn(key), timeout)
        if st != 200:
            continue
        try:
            d = json.loads(body)
        except Exception:
            continue
        info = {'provider': name, 'models': []}
        if name == 'openrouter':
            data = d.get('data', {})
            lim, usage = data.get('limit'), data.get('usage', 0)
            info['balance'] = 'unlimited free-tier' if lim is None else f'${max(0, lim - usage):.2f}/${lim}'
            info['models'] = [data.get('label', 'router')]
        elif name == 'deepseek':
            b = (d.get('balance_infos') or [{}])[0]
            info['balance'] = f"${b.get('total_balance', '?')}"
            info['models'] = ['deepseek-chat', 'deepseek-reasoner']
        elif name == 'moonshot':
            info['balance'] = f"¥{d.get('data', {}).get('balance', '?')}"
        else:
            info['models'] = _parse_models(d)
        alive.append(info)
    return alive


# ── Non-/models probes (different auth shapes) ─────────────────────────────
SPECIAL_PROBES = [
    # (provider, url_template with {key}, header_fn, ok_check)
    ('telegram_bot', 'https://api.telegram.org/bot{key}/getMe', lambda: {}, lambda st, d: st == 200 and d.get('ok')),
    ('huggingface', 'https://huggingface.co/api/whoami-v2', lambda k: {'Authorization': f'Bearer {k}'}, lambda st, d: st == 200),
    ('github_pat', 'https://api.github.com/user', lambda k: {'Authorization': f'token {k}', 'User-Agent': 'kf'}, lambda st, d: st == 200),
    ('npm', 'https://registry.npmjs.org/-/user/org.couchdb.user:me', lambda k: {'Authorization': f'Bearer {k}'}, lambda st, d: st == 200),
    ('slack', 'https://slack.com/api/auth.test', lambda k: {'Authorization': f'Bearer {k}'}, lambda st, d: st == 200 and d.get('ok')),
    ('replicate', 'https://api.replicate.com/v1/account', lambda k: {'Authorization': f'Token {k}'}, lambda st, d: st == 200),
    ('gemini', 'https://generativelanguage.googleapis.com/v1beta/models?key=' + chr(123) + 'key' + chr(125), lambda: {}, lambda st, d: st == 200),
]


def probe_special(key):
    """Probe providers that don't use the OpenAI-style /models endpoint."""
    alive = []
    for name, url_t, hdr_fn, ok in SPECIAL_PROBES:
        try:
            url = url_t.replace('{key}', key)
            h = dict(UA)
            h.update(hdr_fn(key) if hdr_fn.__code__.co_argcount else hdr_fn())
            req = urllib.request.Request(url, headers=h)
            with _OPENER.open(req, timeout=10) as r:
                st = r.status
                try:
                    d = json.loads(r.read())
                except Exception:
                    d = {}
            if ok(st, d):
                info = {'provider': name}
                if name == 'telegram_bot' and d.get('result'):
                    info['models'] = ['@' + d['result'].get('username', '?')]
                if name == 'github_pat' and d.get('login'):
                    info['models'] = [d['login']]
                if name == 'replicate' and d.get('username'):
                    info['models'] = [d['username']]
                alive.append(info)
        except Exception:
            continue
    return alive


# ── one-api / new-api relay probes (D1, ported from jotaro/coff0xc) ──────────
# Relays are self-hosted OpenAI-compatible gateways; their tokens share the
# sk-<hex> shape with DeepSeek/DashScope, so keys DEAD on every known provider
# are often relay tokens. Probing needs the relay base URL — key_patterns
# scan_text captures those as ENDPOINT_CN (ONE_API_URL=... etc).


def _is_local_base(base):
    """True for localhost/private-LAN bases. Leaked .env URLs often point at
    the leaker's LAN — probing those would be SSRF into our own network."""
    try:
        host = (urllib.parse.urlsplit(base).hostname or '').lower()
    except ValueError:
        return True
    if not host or host in ('localhost', '0.0.0.0', '::1') \
            or host.endswith(('.local', '.lan', '.internal')):
        return True
    parts = host.split('.')
    if len(parts) == 4 and all(p.isdigit() for p in parts):
        a, b = int(parts[0]), int(parts[1])
        if a in (0, 10, 127) or a >= 224 or (a, b) in ((169, 254), (192, 168)) \
                or (a == 172 and 16 <= b <= 31):
            return True
    return False


# Auth-gated billing paths: a 200 here proves the key is LIVE on the relay.
RELAY_AUTH_PATHS = (
    ('/api/user/self', ('quota', 'data.quota', 'used_quota', 'data.used_quota')),
    ('/api/user/info', ('quota', 'data.quota', 'used_quota', 'balance')),
    ('/user/info', ('balance', 'quota', 'credits', 'remaining')),
    ('/dashboard/billing/subscription', ('hard_limit_usd', 'system_hard_limit_usd', 'soft_limit_usd')),
    ('/v1/dashboard/billing/subscription', ('hard_limit_usd', 'soft_limit_usd')),
)
# Public path — identifies the relay but NEVER proves the key (no auth needed).
RELAY_ID = ('/api/status', ('quota', 'balance', 'credits'))


def _relay_field(d, fields):
    """First non-zero float among dotted field paths ('data.quota'), else 0.0."""
    for f in fields:
        v = d
        for part in f.split('.'):
            v = v.get(part) if isinstance(v, dict) else None
            if v is None:
                break
        try:
            v = float(v)
        except (TypeError, ValueError):
            continue
        if v:
            # one-api quota unit == USD * 100000 ("500000" means $5)
            return v / 100000 if v > 10000 else v
    return 0.0


def probe_relay(key, base, timeout=3):
    """Probe a one-api/new-api relay at `base` with `key`.

    Returns {'provider': 'relay', 'base', 'live': True, 'balance': usd,
    'source': path} when an auth-gated path answers 200; live=False when only
    the public /api/status answers (relay identified, key unproven); None when
    base is not a responding relay (or is official OpenAI / a local address).
    """
    if (not base or not base.startswith(('http://', 'https://'))
            or 'api.openai.com' in base or _is_local_base(base)):
        return None
    base = base.rstrip('/')
    hdr = {'Authorization': 'Bea' + f'rer {key}'}
    for path, fields in RELAY_AUTH_PATHS:
        st, body = _get(base + path, hdr, timeout)
        if st != 200:
            continue
        try:
            d = json.loads(body)
        except Exception:
            continue
        if isinstance(d, dict):
            return {'provider': 'relay', 'base': base, 'live': True,
                    'balance': _relay_field(d, fields), 'source': path}
    st, body = _get(base + RELAY_ID[0], hdr, timeout)
    if st == 200:
        try:
            d = json.loads(body)
        except Exception:
            return None
        if isinstance(d, dict) and (d.get('success') or 'data' in d):
            return {'provider': 'relay', 'base': base, 'live': False,
                    'balance': _relay_field(d, RELAY_ID[1]), 'source': RELAY_ID[0]}
    return None


def probe_all(key):
    """Full classification: /models endpoints + special probes."""
    return probe(key) + probe_special(key)

def probe_many(keys, workers=20):
    out = {}
    with cf.ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {ex.submit(probe, k): k for k in keys}
        for f in cf.as_completed(futs):
            k = futs[f]
            try:
                r = f.result()
            except Exception:
                r = []
            out[k] = r
    return out


if __name__ == '__main__':
    path = sys.argv[1] if len(sys.argv) > 1 else 'keys.txt'
    relay_base = sys.argv[2] if len(sys.argv) > 2 else ''
    keys = [l.strip() for l in open(path, encoding='utf-8') if l.strip() and not l.startswith('#')]
    print(f'probing {len(keys)} keys against {len(ENDPOINTS)} providers...')
    res = probe_many(keys)
    alive_n = 0
    for k, providers in res.items():
        if providers:
            alive_n += 1
            for p in providers:
                print(f"ALIVE {k[:16]}... -> {p['provider']}: {', '.join(p['models'][:6])}")
    print(f'\ntotal alive: {alive_n}/{len(keys)}')
    if relay_base:
        print(f'\nrelay pass vs {relay_base}')
        for k in keys:
            rr = probe_relay(k, relay_base)
            if rr and rr.get('live'):
                alive_n += 1
                print(f"RELAY {k[:16]}... ${rr['balance']:.2f} via {rr['source']}")