#!/usr/bin/env python3
# The VNC console proxy: nginx passes /novncd/ websockets to it; each one
# carries a console token and is joined to that VM's VNC server, directly,
# through an SSH tunnel or a local Unix socket.

import argparse
import logging
import os
import socket
import sys
from urllib.parse import parse_qs, urlparse

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()

from websockify import ProxyRequestHandler, WebSocketProxy  # noqa: E402

from console.sshtunnels import open_tunnel, tunnel_target  # noqa: E402
from vrtManager.connection import CONN_SOCKET, CONN_SSH  # noqa: E402
from vrtManager.util import is_loopback  # noqa: E402
from webvirtcloud.settings import WS_HOST, WS_PORT  # noqa: E402


def get_parser(default_host=WS_HOST, default_port=WS_PORT or 6080):
    parser = argparse.ArgumentParser(description="WebVirtCloud noVNC 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


def get_connection_infos(token):
    from console.tokens import authorize
    from vrtManager.instance import wvmInstance

    instance, _user = authorize(token)
    conn = wvmInstance(
        instance.compute.hostname,
        instance.compute.login,
        instance.compute.password,
        instance.compute.type,
        instance.name,
        uuid=instance.uuid,  # the VM access was checked for, whatever its name now
    )
    console_host = conn.get_console_listener_addr()
    console_port = conn.get_console_port()
    console_socket = conn.get_console_socket()
    compute = instance.compute
    return compute.hostname, compute.login, compute.type, console_host, console_port, console_socket


WILDCARD = (None, "", "0.0.0.0", "::")


class ConsoleClient:
    """What a console websocket is joined to; ConsoleProxyHandler is the
    websockify handler that runs it."""

    def _new_client(self, socket_factory):
        # The page puts its own console's token on its websocket URL: a cookie
        # is shared by the browser's tabs and would name the last console opened
        token = parse_qs(urlparse(self.path).query).get("token", [None])[0]
        if not token:
            logging.warning("Console connection without a token in its websocket URL")
            return

        try:
            connhost, connuser, conntype, console_host, console_port, console_socket = get_connection_infos(token)
        except Exception as e:
            logging.warning("Console connection denied: %s", e)
            return
        logging.debug(
            "Console of %s@%s (type %s): %s:%s, socket %s",
            connuser, connhost, conntype, console_host, console_port, console_socket,
        )

        tunnel = None
        loopback = is_loopback(console_host)
        if conntype == CONN_SSH:
            # the console travels inside the SSH connection, encrypted,
            # whatever address it listens on
            target = tunnel_target(console_host, console_port, console_socket)
            self.msg("Opening an SSH tunnel to %s@%s, console %s" % (connuser, connhost, target))
            tsock, tunnel = open_tunnel(connhost, connuser, target)
        elif console_socket and conntype == CONN_SOCKET:
            # Local socket on local host
            self.msg("Try to open local socket %s" % console_socket)
            tsock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            tsock.connect(console_socket)
        elif conntype == CONN_SOCKET and loopback:
            # libvirt runs on this host: its loopback is ours
            self.msg("connecting to: %s:%s" % (console_host, console_port))
            tsock = socket_factory(console_host, console_port, connect=True)
        elif console_socket or loopback:
            message = "The console listens only on its compute; it is reached through SSH computes only"
            self.msg(message)
            raise Exception(message)
        else:
            # the address the console listens on, or the compute's for a wildcard
            address = connhost if console_host in WILDCARD else console_host
            self.msg("connecting to: %s:%s" % (address, console_port))
            tsock = socket_factory(address, console_port, connect=True)

        try:
            self.msg("Start proxying")
            self.do_proxy(tsock)
        finally:
            tsock.close()
            if tunnel:
                error = tunnel.close()
                if error:
                    logging.warning("SSH tunnel to %s@%s failed: %s", connuser, connhost, error)


class ConsoleProxyHandler(ProxyRequestHandler, ConsoleClient):
    def msg(self, *args, **kwargs):
        self.log_message(*args, **kwargs)

    def new_websocket_client(self):
        self._new_client(self.server.socket)


if __name__ == "__main__":
    options = get_parser().parse_args()

    # to stderr (the journal, docker logs), not the panel's own log file:
    # each process rotating that file would lose lines. force: Django's
    # settings already gave the root logger that file.
    level = logging.DEBUG if options.debug else logging.INFO if options.verbose else logging.WARNING
    logging.basicConfig(level=level, format="%(levelname)s %(name)s: %(message)s", force=True)
    options.verbose = options.verbose or options.debug

    # nginx ends TLS: novncd speaks plain websockets on the loopback
    WebSocketProxy(
        RequestHandlerClass=ConsoleProxyHandler,
        listen_host=options.host,
        listen_port=options.port,
        verbose=options.verbose,
        daemon=False,
        web=False,
        target_host="ignore",
        target_port="ignore",
        # a ping while the VM's screen stays the same: nginx, and a proxy on
        # the way, end a websocket that sends nothing for 60 seconds
        heartbeat=30,
    ).start_server()