backnforth/server.py

"""backnforth server — static site + WebSocket relay, Python stdlib only.

    python3 server.py [--no-browser]

Serves the therapist page on http://<ip>:8000/ and the session relay on
ws://<ip>:8787. No pip installs, no virtualenv — any Python 3.7+ works,
including old machines that will never see a package index.
"""
import base64
import hashlib
import json
import os
import socket
import socketserver
import struct
import sys
import threading
import webbrowser
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer

HTTP_PORT = 8000
WS_PORT = 8787
_WS_GUID = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

SESSIONS = {}   # session id -> set of live WSHandler
LOCK = threading.Lock()
_SERVERS = []
TOKENS = None   # None = open relay; a set = hosts must present a valid access code


def local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except OSError:
        try:
            return socket.gethostbyname(socket.gethostname())
        except OSError:
            return '127.0.0.1'


def shutdown_all():
    for srv in _SERVERS:
        try:
            srv.shutdown()
        except Exception:
            pass


class WSHandler(socketserver.BaseRequestHandler):
    """Minimal RFC6455 server: unfragmented text frames, ping/pong, close.

    ponytail: no fragmentation/binary support — the state payloads are a few
    hundred bytes of JSON and browsers never fragment at that size.
    """

    def setup(self):
        self.session_id = None
        self.send_lock = threading.Lock()

    def handle(self):
        if not self._handshake():
            return
        try:
            while True:
                msg = self._read_message()
                if msg is None:
                    break
                self._dispatch(msg)
        except (ConnectionError, OSError):
            pass

    def finish(self):
        if self.session_id:
            with LOCK:
                SESSIONS.get(self.session_id, set()).discard(self)

    # ---- relay protocol (same JSON messages as the old websockets server) ----

    def _dispatch(self, text):
        try:
            data = json.loads(text)
        except ValueError:
            return
        t = data.get('type')
        if t == 'join':
            sid = data.get('sessionId') or 'default'
            if TOKENS is not None:
                # hosts need an access code; patients may only join a LIVE session
                # (they got the random room id from their therapist's link)
                if data.get('role') == 'patient':
                    with LOCK:
                        live = bool(SESSIONS.get(sid))
                    if not live:
                        self._send_json({'type': 'join_denied', 'reason': 'session not live'})
                        return
                elif data.get('token') not in TOKENS:
                    self._send_json({'type': 'join_denied', 'reason': 'access code required'})
                    return
            self.session_id = sid
            with LOCK:
                SESSIONS.setdefault(self.session_id, set()).add(self)
            self._send_json({'type': 'joined', 'sessionId': self.session_id})
        elif t in ('state', 'patientViewport'):
            if TOKENS is not None and not self.session_id:
                return  # hardened relay only forwards for joined connections
            sid = (self.session_id if TOKENS is not None
                   else data.get('sessionId') or self.session_id or 'default')
            out = json.dumps({'type': t, 'payload': data.get('payload')})
            with LOCK:
                peers = list(SESSIONS.get(sid, ()))
            for peer in peers:
                if peer is not self:
                    try:
                        peer._send_text(out)
                    except Exception:
                        pass
        elif t in ('whoami', 'detect_ip'):
            self._send_json({'type': 'detect_ip', 'ip': local_ip()})
        elif t == 'ping':
            self._send_json({'type': 'pong'})
        elif t == 'shutdown':
            if self.client_address[0] in ('127.0.0.1', '::1', local_ip()):
                self._send_json({'type': 'shutdown_ack'})
                threading.Thread(target=shutdown_all, daemon=True).start()

    # ---- websocket plumbing ----

    def _handshake(self):
        self.request.settimeout(10)
        data = b''
        while b'\r\n\r\n' not in data:
            chunk = self.request.recv(4096)
            if not chunk or len(data) > 65536:
                return False
            data += chunk
        headers = {}
        for line in data.split(b'\r\n')[1:]:
            if b':' in line:
                k, v = line.split(b':', 1)
                headers[k.strip().lower()] = v.strip()
        key = headers.get(b'sec-websocket-key')
        if not key:
            return False
        accept = base64.b64encode(hashlib.sha1(key + _WS_GUID).digest()).decode()
        self.request.sendall((
            'HTTP/1.1 101 Switching Protocols\r\n'
            'Upgrade: websocket\r\n'
            'Connection: Upgrade\r\n'
            'Sec-WebSocket-Accept: ' + accept + '\r\n\r\n'
        ).encode())
        self.request.settimeout(None)
        return True

    def _recv_exact(self, n):
        buf = b''
        while len(buf) < n:
            chunk = self.request.recv(n - len(buf))
            if not chunk:
                return None
            buf += chunk
        return buf

    def _read_message(self):
        while True:
            head = self._recv_exact(2)
            if head is None:
                return None
            b1, b2 = head
            opcode = b1 & 0x0F
            length = b2 & 0x7F
            if length == 126:
                ext = self._recv_exact(2)
                if ext is None:
                    return None
                length = struct.unpack('>H', ext)[0]
            elif length == 127:
                ext = self._recv_exact(8)
                if ext is None:
                    return None
                length = struct.unpack('>Q', ext)[0]
            if length > 1_000_000:
                return None
            mask = self._recv_exact(4) if (b2 & 0x80) else b'\x00\x00\x00\x00'
            if mask is None:
                return None
            payload = self._recv_exact(length) if length else b''
            if payload is None:
                return None
            payload = bytes(c ^ mask[i % 4] for i, c in enumerate(payload))
            if opcode == 0x1:    # text
                try:
                    return payload.decode('utf-8')
                except UnicodeDecodeError:
                    return None
            elif opcode == 0x9:  # ping -> pong
                self._send_frame(0xA, payload)
            elif opcode == 0x8:  # close
                try:
                    self._send_frame(0x8, b'')
                except Exception:
                    pass
                return None
            # anything else (binary/continuation): skip

    def _send_frame(self, opcode, payload):
        n = len(payload)
        head = bytes([0x80 | opcode])
        if n < 126:
            head += bytes([n])
        elif n < 65536:
            head += bytes([126]) + struct.pack('>H', n)
        else:
            head += bytes([127]) + struct.pack('>Q', n)
        with self.send_lock:
            self.request.sendall(head + payload)

    def _send_text(self, text):
        self._send_frame(0x1, text.encode('utf-8'))

    def _send_json(self, obj):
        self._send_text(json.dumps(obj))


class _QuietHTTP(SimpleHTTPRequestHandler):
    def log_message(self, *args):
        pass


def _argval(flag, default):
    args = sys.argv[1:]
    return args[args.index(flag) + 1] if flag in args else default


def main():
    global TOKENS
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    bind = _argval('--bind', '0.0.0.0')          # e.g. a tailnet IP behind a reverse proxy
    http_port = int(_argval('--http-port', HTTP_PORT))
    tokens_file = _argval('--tokens', None)      # one access code per line, # comments ok
    if tokens_file:
        with open(tokens_file, encoding='utf-8') as f:
            TOKENS = {l.strip() for l in f if l.strip() and not l.startswith('#')}
        print('relay hardened: %d access codes loaded' % len(TOKENS))

    socketserver.ThreadingTCPServer.allow_reuse_address = True
    ThreadingHTTPServer.allow_reuse_address = True
    httpd = ThreadingHTTPServer((bind, http_port), _QuietHTTP)
    wsd = socketserver.ThreadingTCPServer((bind, WS_PORT), WSHandler)
    httpd.daemon_threads = True
    wsd.daemon_threads = True
    _SERVERS.extend([httpd, wsd])

    ip = bind if bind not in ('0.0.0.0', '') else local_ip()
    print('backnforth running:')
    print('  Therapist page:  http://%s:%d/' % (ip, http_port))
    print('  Session relay:   ws://%s:%d' % (ip, WS_PORT))
    print('Press Ctrl+C to stop.')

    threading.Thread(target=wsd.serve_forever, daemon=True).start()
    if '--no-browser' not in sys.argv:
        try:
            webbrowser.open('http://%s:%d/easy.html' % (ip, http_port))
        except Exception:
            pass
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print('\nStopped.')


if __name__ == '__main__':
    main()