webvirtcloud/console/socketiod

#!/usr/bin/env python3
import os
import sys
import logging
import django

DIR_PATH = os.path.dirname(os.path.abspath(__file__))
ROOT_PATH = os.path.abspath(os.path.join(DIR_PATH, "..", ""))
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "webvirtcloud.settings")
CERT = DIR_PATH + "/cert.pem"

if ROOT_PATH not in sys.path:
    sys.path.append(ROOT_PATH)

django.setup()

from django.conf import settings as django_settings

# socketiod does not authenticate connections yet: any client
# that can reach the port can read and write an open serial console. Refuse
# to start unless an admin explicitly opts in.
if not getattr(django_settings, "SERIAL_CONSOLE_ENABLED", False):
    sys.stderr.write(
        "socketiod: the serial console is disabled (SERIAL_CONSOLE_ENABLED is not "
        "True in settings). It is unauthenticated; only enable it on a trusted "
        "network.\n"
    )
    sys.exit(1)

import codecs
import pty
import select
import struct
import fcntl
import termios
import signal
import socketio
import eventlet

# from six.moves import http_cookies as Cookie
from http import cookies as Cookie
from webvirtcloud.settings import SOCKETIO_PORT, SOCKETIO_HOST
import argparse


def get_parser(default_host=SOCKETIO_HOST, default_port=SOCKETIO_PORT or 6081):
    parser = argparse.ArgumentParser(description="WebVirtCloud Socket.IO daemon")
    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbose",
        action="store_true",
        help="Verbose mode",
        default=False,
    )
    parser.add_argument(
        "-d", "--debug", dest="debug", action="store_true", help="Debug mode", default=False
    )
    parser.add_argument(
        "-H",
        "--host",
        dest="host",
        action="store",
        help="Listen host",
        default=default_host,
    )
    parser.add_argument(
        "-p",
        "--port",
        dest="port",
        action="store",
        type=int,
        help="Listen port",
        default=default_port,
    )
    return parser


options = get_parser().parse_args()

FORMAT = "%(asctime)s - %(name)s - %(levelname)s : %(message)s"
if options.debug:
    logging.basicConfig(level=logging.DEBUG, format=FORMAT)
    options.verbose = True
elif options.verbose:
    logging.basicConfig(level=logging.INFO, format=FORMAT)
else:
    logging.basicConfig(level=logging.WARNING, format=FORMAT)

async_mode = "eventlet"
sio = socketio.Server(async_mode=async_mode, cors_allowed_origins=[])

fd = None
child_pid = None


def get_connection_infos(token):
    from instances.models import Instance
    from vrtManager.instance import wvmInstance

    try:
        temptoken = token.split("-", 1)
        host = int(temptoken[0])
        uuid = temptoken[1]
        instance = Instance.objects.get(compute_id=host, uuid=uuid)
        conn = wvmInstance(
            instance.compute.hostname,
            instance.compute.login,
            instance.compute.password,
            instance.compute.type,
            instance.name,
        )
    except Exception as e:
        logging.error(
            "Fail to retrieve console connection infos for token %s : %s" % (token, e)
        )
        raise
    return (instance, conn)


def set_winsize(fd, row, col, xpix=0, ypix=0):
    winsize = struct.pack("HHHH", row, col, xpix, ypix)
    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)


def stop_console():
    """Stop the console process (if still running), reap it and close its pty."""
    global fd
    global child_pid

    if child_pid:
        try:
            os.kill(child_pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        try:
            os.waitpid(child_pid, 0)  # only the child this daemon started
        except ChildProcessError:
            pass
    if fd is not None:
        os.close(fd)
    fd = None
    child_pid = None


def read_and_forward_pty_output():
    max_read_bytes = 1024 * 20
    # a character may be split between two reads
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    # The pty number may be reused by the next console; the pid identifies
    # this one. Every yield (sleep, emit) may end it.
    console_pid, console_fd = child_pid, fd
    while True:
        sio.sleep(0.01)
        if child_pid != console_pid:
            return
        (data_ready, _, _) = select.select([console_fd], [], [], 0)
        if not data_ready:
            continue
        try:
            output = os.read(console_fd, max_read_bytes)
        except OSError:  # EIO: the console process ended
            output = b""
        if not output:
            sio.emit("pty_output", {"output": decoder.decode(b"", final=True) + "\r\n[console closed]\r\n"})
            stop_console()
            for sid, _ in list(sio.manager.get_participants("/", None)):
                sio.disconnect(sid)
            return
        sio.emit("pty_output", {"output": decoder.decode(output)})


@sio.event
def resize(sid, message):
    global fd
    if fd:
        set_winsize(fd, message["rows"], message["cols"])


@sio.event
def pty_input(sid, message):
    global fd
    if fd:
        os.write(fd, message["input"].encode())


@sio.event
def disconnect_request(sid):
    sio.disconnect(sid)


@sio.event
def connect(sid, environ):
    global fd
    global child_pid

    token = None
    hcookie = environ.get("HTTP_COOKIE")
    if hcookie:
        cookie = Cookie.SimpleCookie()
        for hcookie_part in hcookie.split(";"):
            hcookie_part = hcookie_part.lstrip()
            try:
                cookie.load(hcookie_part)
            except Cookie.CookieError:
                logging.warning("Found malformed cookie")
            else:
                if "token" in cookie:
                    token = cookie["token"].value

    if child_pid:
        # already started child process, don't start another
        # write a new line so that when a client refresh the shell prompt is printed
        os.write(fd, str.encode("\n"))
        return

    # create child process attached to a pty we can read from and write to
    (child_pid, fd) = pty.fork()

    if child_pid == 0:
        # Become the console process; never return, which would run this
        # server's code in the child as well. Its output goes to the client.
        try:
            if token is None:
                sys.stderr.write("No console token\n")
            else:
                (instance, conn) = get_connection_infos(token)
                os.execv("conf/daemon/consolecallback", ["consolecallback", conn.wvm.getURI(), conn.get_uuid()])
        except BaseException:
            pass  # get_connection_infos logged it
        finally:
            os._exit(1)
    else:
        # this is the parent process fork.
        sio.start_background_task(target=read_and_forward_pty_output)


@sio.event
def disconnect(sid):
    stop_console()


app = socketio.WSGIApp(sio)

eventlet.wsgi.server(eventlet.listen((options.host, int(options.port))), app)