#!/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()

import re
import socket

from urllib.parse import parse_qs, urlparse
import argparse
from console.sshtunnels import close_tunnel, open_tunnel, tunnel_target
from vrtManager.connection import CONN_SOCKET, CONN_SSH
from webvirtcloud.settings import WS_CERT, WS_HOST, WS_PORT


def get_parser(default_host=WS_HOST, default_port=WS_PORT or 6080, default_cert=WS_CERT or CERT):
    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,
    )
    parser.add_argument(
        "-c",
        "--cert",
        dest="cert",
        action="store",
        help="Certificate file path",
        default=default_cert,
    )
    return parser



try:
    from websockify import WebSocketProxy

    try:
        from websockify import ProxyRequestHandler
    except ImportError:
        USE_HANDLER = False
    else:
        USE_HANDLER = True
except ImportError:
    try:
        from novnc.wsproxy import WebSocketProxy
    except ImportError:
        print("Unable to import a websockify implementation,\n please install one")
        sys.exit(1)
    else:
        USE_HANDLER = False


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

    try:
        instance = 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()
    except Exception as e:
        logging.error("Fail to retrieve console connection infos: %s", e)
        raise
    compute = instance.compute
    return compute.hostname, compute.login, compute.type, console_host, console_port, console_socket


class CompatibilityMixIn(object):
    def _new_client(self, daemon, 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:
            self.msg("No console token in the websocket URL")
            return

        try:
            connhost, connuser, conntype, console_host, console_port, console_socket = get_connection_infos(token)
        except Exception as e:
            self.msg("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 = re.match(r"^127\.", console_host or "")
        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:
            # Direct access
            self.msg("connecting to: %s:%s" % (connhost, console_port))
            tsock = socket_factory(connhost, console_port, connect=True)

        if self.verbose and not daemon:
            print(self.traffic_legend)

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


if USE_HANDLER:

    class NovaProxyRequestHandler(ProxyRequestHandler, CompatibilityMixIn):
        def msg(self, *args, **kwargs):
            self.log_message(*args, **kwargs)

        def vmsg(self, *args, **kwargs):
            if self.verbose:
                self.msg(*args, **kwargs)

        def new_websocket_client(self):
            """
            Called after a new WebSocket connection has been established.
            """
            # Setup variable for compatibility
            daemon = self.server.daemon
            socket_factory = self.server.socket

            self._new_client(daemon, socket_factory)

else:

    class NovaWebSocketProxy(WebSocketProxy, CompatibilityMixIn):
        def new_client(self):
            """
            Called after a new WebSocket connection has been established.
            """
            # Setup variable for compatibility
            daemon = self.daemon
            socket_factory = self.socket

            self._new_client(daemon, socket_factory)


if __name__ == "__main__":
    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)

    if USE_HANDLER:
        # Create the WebSocketProxy with NovaProxyRequestHandler handler
        server = WebSocketProxy(
            RequestHandlerClass=NovaProxyRequestHandler,
            listen_host=options.host,
            listen_port=options.port,
            source_is_ipv6=False,
            verbose=options.verbose,
            cert=options.cert,
            key=None,
            ssl_only=False,
            daemon=False,
            record=False,
            web=False,
            traffic=False,
            target_host="ignore",
            target_port="ignore",
            wrap_mode="exit",
            wrap_cmd=None,
        )
    else:
        # Create the NovaWebSockets proxy
        server = NovaWebSocketProxy(
            listen_host=options.host,
            listen_port=options.port,
            source_is_ipv6=False,
            verbose=options.verbose,
            cert=options.cert,
            key=None,
            ssl_only=False,
            daemon=False,
            record=False,
            web=False,
            target_host="ignore",
            target_port="ignore",
            wrap_mode="exit",
            wrap_cmd=None,
        )
    server.start_server()
