webvirtcloud/console/socketiod

#!/usr/bin/env python3
# The serial console daemon. The page connects with its signed console token
# (Socket.IO auth); each accepted connection gets its own console worker
# (conf/daemon/consolecallback) in a pty, which opens the VM's console with the
# compute's credentials. This process never opens libvirt. Connections are
# refused while SERIAL_CONSOLE_ENABLED is off.
import argparse
import codecs
import fcntl
import os
import pty
import select
import signal
import struct
import sys
import termios

import django

ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
if ROOT_PATH not in sys.path:
    sys.path.append(ROOT_PATH)
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "webvirtcloud.settings")
django.setup()

import socketio  # noqa: E402
from gunicorn.app.base import BaseApplication  # noqa: E402
from django import db  # noqa: E402
from django.conf import settings  # noqa: E402

from console.tokens import authorize  # noqa: E402

CONSOLECALLBACK = os.path.join(ROOT_PATH, "conf", "daemon", "consolecallback")
MAX_SESSIONS = 20

sio = socketio.Server(async_mode="threading")  # the same origin only
sessions = {}  # sid -> {"pid", "fd", "vm"}


@sio.event
def connect(sid, environ, auth):
    if not settings.SERIAL_CONSOLE_ENABLED:
        return False
    try:
        instance = authorize((auth or {}).get("token"))
    except PermissionError:
        return False
    finally:
        db.close_old_connections()
    if len(sessions) >= MAX_SESSIONS or any(s["vm"] == instance.pk for s in sessions.values()):
        return False
    pid, fd = pty.fork()
    if pid == 0:
        try:
            os.execv(sys.executable, [sys.executable, CONSOLECALLBACK, str(instance.pk)])
        finally:
            os._exit(1)
    sessions[sid] = session = {"pid": pid, "fd": fd, "vm": instance.pk}
    sio.start_background_task(forward, sid, session)


def forward(sid, session):
    """The worker's output to its client, until the worker ends."""
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")  # a character may be split
    while sessions.get(sid) is session:
        if not select.select([session["fd"]], [], [], 0.1)[0]:
            continue
        try:
            output = os.read(session["fd"], 20480)
        except OSError:  # EIO: the worker ended
            output = b""
        if not output:
            sio.emit("pty_output", {"output": decoder.decode(b"", final=True) + "\r\n[console closed]\r\n"}, to=sid)
            stop(sid)
            sio.disconnect(sid)
            return
        sio.emit("pty_output", {"output": decoder.decode(output)}, to=sid)


def stop(sid):
    """End the sid's worker; once, whoever calls it first."""
    session = sessions.pop(sid, None)
    if session is None:
        return
    try:
        os.close(session["fd"])
    except OSError:
        pass
    try:
        os.kill(session["pid"], signal.SIGKILL)
    except ProcessLookupError:
        pass
    sio.start_background_task(os.waitpid, session["pid"], 0)


@sio.event
def pty_input(sid, message):
    session = sessions.get(sid)
    if session and isinstance(message, dict) and isinstance(message.get("input"), str):
        try:
            os.write(session["fd"], message["input"].encode())
        except OSError:
            pass


@sio.event
def resize(sid, message):
    session = sessions.get(sid)
    if not (session and isinstance(message, dict)):
        return
    rows, cols = message.get("rows"), message.get("cols")
    if all(isinstance(v, int) and 1 <= v <= 500 for v in (rows, cols)):
        fcntl.ioctl(session["fd"], termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))


@sio.event
def disconnect(sid, reason=None):
    stop(sid)


class Server(BaseApplication):
    """gunicorn with one worker (the sessions live in it); a thread per connection."""

    def __init__(self, host, port):
        self.bind = "%s:%d" % (host, port)
        super().__init__()

    def load_config(self):
        self.cfg.set("bind", self.bind)
        self.cfg.set("worker_class", "gthread")
        self.cfg.set("threads", 2 * MAX_SESSIONS)
        self.cfg.set("graceful_timeout", 1)  # open consoles never finish; their workers end with the pty
        self.cfg.set("control_socket_disable", True)

    def load(self):
        return socketio.WSGIApp(sio)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebVirtCloud serial console daemon")
    parser.add_argument("-H", "--host", default=settings.SOCKETIO_HOST, help="Listen host")
    parser.add_argument("-p", "--port", type=int, default=settings.SOCKETIO_PORT, help="Listen port")
    options = parser.parse_args()
    Server(options.host, options.port).run()