# -*- coding: utf-8 -*-
"""localtunnel.me 客户端（Python 移植，协议见 localtunnel/localtunnel 源码）。

- 控制面：GET https://localtunnel.me/<subdomain> (Accept: json) 预留固定子域名，
  返回 {id, port, max_conn_count, url}；地址形如 https://<sub>.loca.lt，不随时间轮换。
- 数据面：max_conn 条裸 TCP 长连接到 localtunnel.me:<port>；每条连接承载一次公网
  HTTP 请求的原始字节，本地 pipe 到 127.0.0.1:8899；连接死了就重连，URL 不变。
- 成功拿到 url 后写回 tunnel_url.txt（仅变化时写）。

用法: nohup python3 lt_tunnel.py [subdomain] > /tmp/lt_tunnel.log 2>&1 &
"""
import json, os, selectors, socket, sys, time, urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
SUB = sys.argv[1] if len(sys.argv) > 1 else 'south-ops-hb'
LOCAL_PORT = 8899
CTRL_HOST = 'localtunnel.me'
URL_FILE = os.path.join(HERE, 'tunnel_url.txt')

def log(*a):
    print(time.strftime('%m-%d %H:%M:%S'), *a, flush=True)

def claim(sub):
    req = urllib.request.Request('https://%s/%s' % (CTRL_HOST, sub),
                                 headers={'Accept': 'application/json'})
    with urllib.request.urlopen(req, timeout=15) as r:
        return json.loads(r.read().decode())

def pipe_pair(r, l):
    sel = selectors.DefaultSelector()
    sel.register(r, selectors.EVENT_READ, r)
    sel.register(l, selectors.EVENT_READ, l)
    try:
        while True:
            for key, _ in sel.select(timeout=120):
                data = key.fileobj.recv(1 << 16)
                if not data:
                    return
                (l if key.fileobj is r else r).sendall(data)
    except OSError:
        pass
    finally:
        sel.close()

def slot(port, idx):
    while True:
        try:
            r = socket.create_connection((CTRL_HOST, port), timeout=30)
            r.settimeout(None)
            r.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            log('slot%d tunnel connected' % idx)
            while True:
                first = r.recv(1 << 16)          # 等公网请求被分配到这条连接
                if not first:
                    break
                try:
                    l = socket.create_connection(('127.0.0.1', LOCAL_PORT), timeout=10)
                except OSError:
                    log('slot%d local 8899 refused' % idx)
                    break
                try:
                    l.sendall(first)
                    pipe_pair(r, l)
                finally:
                    l.close()
        except OSError as e:
            log('slot%d err: %s' % (idx, e))
        finally:
            try:
                r.close()
            except Exception:
                pass
        time.sleep(2)

def main():
    while True:
        try:
            info = claim(SUB)
            url, port = info.get('url'), info.get('port')
            n = max(1, int(info.get('max_conn_count') or 1))
            if not url or not port:
                raise ValueError(info)
        except Exception as e:
            log('claim failed:', e, '- retry 10s')
            time.sleep(10)
            continue
        log('claimed', url, 'port', port, 'slots', n)
        try:
            old = open(URL_FILE).read().strip()
        except OSError:
            old = ''
        if old != url:
            open(URL_FILE, 'w').write(url + '\n')
            log('tunnel_url.txt updated:', url)
        for i in range(n):
            os.fork() == 0 and (_fork_child(port, i))
        # 父进程守候：子进程全死则重新 claim（理论上不会）
        while True:
            pid, _ = os.wait()
            log('child', pid, 'exited; respawning slot')
            os.fork() == 0 and _fork_child(port, 0)

def _fork_child(port, idx):
    try:
        slot(port, idx)
    finally:
        os._exit(0)

if __name__ == '__main__':
    main()