#!/usr/bin/env python3
# consolecallback <instance id> <user id> - the serial console of one VM on stdin/stdout
# (socketiod runs it in a pty, one per client). It opens the console through
# the compute's connection (its stored credentials); the connection's event
# loop thread runs the stream callbacks. The console survives guest reboots.

import os
import sys
import time
import tty

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 libvirt  # noqa: E402
from django import db  # noqa: E402
from django.contrib.auth import get_user_model  # noqa: E402

from instances.models import Instance  # noqa: E402
from instances.utils import can_open_console  # noqa: E402
from operations.utils import public_error_text  # noqa: E402
from vrtManager.connection import connection_manager  # noqa: E402

# libvirt prints its errors on stderr, which is the user's terminal: they are
# raised and told in one line instead
libvirt.registerErrorHandler(lambda context, error: None, None)
SEND_WAIT = 5  # seconds a paste may wait for the VM to take it
ACCESS_CHECK = 10  # seconds between checks that the user may still use the console


def say(text):
    os.write(1, ("\r\n[%s]\r\n" % text).encode())


class Console:
    def __init__(self, connection, uuid):
        self.domain = connection.lookupByUUIDString(uuid)
        self.connection = connection
        self.stream = None
        self.running = True

    def check(self):
        """Open the stream while the VM runs, drop it while it does not."""
        active = self.domain.state()[0] in (libvirt.VIR_DOMAIN_RUNNING, libvirt.VIR_DOMAIN_PAUSED)
        if active and self.stream is None:
            stream = self.connection.newStream(libvirt.VIR_STREAM_NONBLOCK)
            self.domain.openConsole(None, stream, 0)  # a console in use is refused
            stream.eventAddCallback(libvirt.VIR_STREAM_EVENT_READABLE, from_vm, self)
            self.stream = stream
        elif not active and self.stream is not None:
            self.close_stream()
            say("The VM is not running: its console opens again when it starts")

    def close_stream(self):
        stream, self.stream = self.stream, None
        if stream is None:  # the other thread closed it
            return
        try:
            stream.eventRemoveCallback()
            stream.abort()
        except libvirt.libvirtError:
            pass


def from_client(watch, fd, events, console):
    try:
        data = os.read(fd, 1024)
    except OSError:
        data = b""
    if not data:  # the client is gone
        console.running = False
        return
    stream = console.stream  # the main thread may drop it meanwhile
    deadline = time.monotonic() + SEND_WAIT
    try:
        # a non-blocking stream takes part of a paste, or none (-2) while the VM is busy
        while data and stream is not None and time.monotonic() < deadline:
            sent = stream.send(data)
            if sent == -2:
                time.sleep(0.01)
            else:
                data = data[sent:]
    except libvirt.libvirtError:
        console.close_stream()


def from_vm(stream, events, console):
    try:
        data = stream.recv(1024)
    except libvirt.libvirtError:  # the VM shut down
        data = None
    if not data or data == -2:  # -2: nothing to read now
        if data != -2:
            console.close_stream()
        return
    os.write(1, data)


def may_still_open(user_id, instance):
    """The user still may open the console: an owner may have been removed or
    the user deactivated since the session started."""
    try:
        user = get_user_model().objects.filter(pk=user_id).first()
        return user is not None and can_open_console(user, instance)
    finally:
        db.connection.close()  # a session may last for days: no idle connection meanwhile


def main():
    instance = Instance.objects.select_related("compute").get(pk=int(sys.argv[1]))
    user_id = int(sys.argv[2])
    c = instance.compute
    try:
        console = Console(connection_manager.get_connection(c.hostname, c.login, c.password, c.type), instance.uuid)
        tty.setraw(0)  # before the console opens: nothing typed is echoed by the terminal
        console.check()
    except libvirt.libvirtError as err:
        say("Cannot open the console of %s: %s" % (instance.name, public_error_text(err)))
        sys.exit(1)
    if console.stream is None:
        say("The VM is not running: its console opens when it starts")
    libvirt.virEventAddHandle(0, libvirt.VIR_EVENT_HANDLE_READABLE, from_client, console)
    checked = time.monotonic()
    try:
        while console.running:
            time.sleep(0.2)
            console.check()
            if time.monotonic() - checked > ACCESS_CHECK:
                if not may_still_open(user_id, instance):
                    say("You may no longer open this console")
                    sys.exit(1)
                checked = time.monotonic()
    except libvirt.libvirtError as err:  # the host went away, the VM was undefined or migrated
        say("The console of %s closed: %s" % (instance.name, public_error_text(err)))
        sys.exit(1)


if __name__ == "__main__":
    main()