open-standards/wai/wai-rs/tests/fixtures/dev-stack/fake-stack.py

#!/usr/bin/env python3
"""A stand-in for the relay and the room service, for tests/cli_studio.rs: what `wai dev`
asks of each, and nothing more.

  fake-stack.py keygen|relay-lines <keys.toml>   prints relay lines (comments here)
  fake-stack.py --config <relay.toml|room.toml>  listens where the configuration says

A relay's configuration names `websocket_listen`; it accepts connections and closes them. A
room service's names `listen`; it answers GET /health as a service whose clock is on the
relay, and GET /room with an empty room. Either exits when the process that started it is
gone. Standard library only.
"""
import http.server
import json
import os
import re
import socket
import socketserver
import sys
import threading
import time


def watch_parent():
    parent = os.getppid()
    while True:
        time.sleep(0.5)
        if os.getppid() != parent:
            os._exit(0)


def listen_addr(config, key):
    text = open(config, encoding="utf-8").read()
    m = re.search(r'^%s\s*=\s*"([^"]+)"' % key, text, re.M)
    host, port = m.group(1).rsplit(":", 1)
    return host, int(port)


class Room(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path.startswith("/health"):
            body = {"ok": True, "relay": True, "room": "00" * 32}
        elif self.path.startswith("/room"):
            body = {"room": "00" * 32, "t0_us": 0, "descriptor": "", "grants": []}
        else:
            self.send_error(404)
            return
        b = json.dumps(body).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def log_message(self, *args):
        pass


class Server(http.server.ThreadingHTTPServer):
    """HTTPServer looks up a name for its address before it listens, and nothing here reads the
    name. On some hosts that lookup takes 35 s, past what `wai dev` waits for a service."""

    def server_bind(self):
        socketserver.TCPServer.server_bind(self)


def main(argv):
    if argv[:1] in (["keygen"], ["relay-lines"]):
        open(argv[1], "a").close()
        print("# authority %s\n# clock     %s" % ("11" * 32, "22" * 32))
        return 0
    if argv[:1] != ["--config"]:
        print("fake-stack.py: keygen|relay-lines <keys> or --config <file>", file=sys.stderr)
        return 2
    threading.Thread(target=watch_parent, daemon=True).start()
    config = argv[1]
    if "websocket_listen" in open(config, encoding="utf-8").read():
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(listen_addr(config, "websocket_listen"))
        s.listen(16)
        while True:
            c, _ = s.accept()
            c.close()
    Server(listen_addr(config, "listen"), Room).serve_forever()
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))