#!/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 ipaddress
import logging
import os
import socket
import sys
import threading
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 console.tokens import ACCESS_CHECK, authorize, may_still_open, session_user_id  # noqa: E402
from logs.audit import write  # noqa: E402
from vrtManager.connection import CONN_SOCKET, CONN_SSH, CONNECT_TIMEOUT  # 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(instance):
    from vrtManager.instance import wvmInstance

    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", "::")


def bounded_connect(host, port, connect=True):
    """A connection to a VNC server, in place of websockify's socket(), which
    waits for one that does not answer as long as the OS does (some 2 minutes)."""
    tsock = socket.create_connection((host, port), timeout=CONNECT_TIMEOUT)
    tsock.settimeout(None)  # a console may stay quiet for long once connected
    return tsock


def client_address(value):
    """The client's address nginx sends in X-Real-IP, or None."""
    try:
        return str(ipaddress.ip_address(value))
    except ValueError:
        return None


def end_when_refused(instance, user_id, tsock, ended):
    """Close the session's VNC socket, which ends it, once its user may no
    longer open the console (as the serial console does)."""
    while not ended.wait(ACCESS_CHECK):
        if not may_still_open(user_id, instance):
            logging.info("VNC console of VM %s closed: its user may no longer open it", instance.name)
            try:
                tsock.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass  # the session ended meanwhile
            return


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

        user = None
        try:
            instance, user = authorize(token, "vnc")
            # a token read from a log or a URL is no use without its user's session
            if session_user_id(self.headers.get("Cookie")) != str(user.pk):
                raise PermissionError("The console page was opened in another session")
            connhost, connuser, conntype, console_host, console_port, console_socket = get_connection_infos(instance)
        except PermissionError as e:
            address = client_address(self.headers.get("X-Real-IP"))
            logging.warning("Console connection from %s denied: %s", address, e)
            write(user.username if user else "", "-", "", f"VNC console refused: {e}", False, address)
            return
        except Exception as e:
            logging.warning("Console connection failed: %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)

        ended = threading.Event()
        threading.Thread(target=end_when_refused, args=(instance, user.pk, tsock, ended), daemon=True).start()
        try:
            self.msg("Start proxying")
            self.do_proxy(tsock)
        finally:
            ended.set()
            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(bounded_connect)


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
    server = 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,
    )
    try:
        server.start_server()
    except SystemExit as stop:
        # websockify ends with status 0 when it cannot listen (the port is
        # taken): systemd's Restart=on-failure would not start it again
        raise SystemExit(stop.code or 1) from None
