#!/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.
# socketiod made the terminal raw before starting it (setting it here would
# drop what the user typed meanwhile).

import logging
import os
import sys
import time

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 instances.models import Instance  # noqa: E402
from console.tokens import ACCESS_CHECK, may_still_open  # noqa: E402
from operations.utils import public_error_text  # noqa: E402
from vrtManager.connection import connection_manager  # noqa: E402

SEND_WAIT = 5  # seconds a paste may wait for the VM to take it


log = logging.getLogger("consolecallback")


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.active = False  # the VM ran at the last check
        self.running = True

    def check(self):
        """Open the stream while the VM runs, drop it while it does not. The
        stream may have ended first: the VM's stop is still reported."""
        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)
            try:
                self.domain.openConsole(None, stream, 0)  # a console in use is refused
            except libvirt.libvirtError:
                if self.domain.isActive():
                    raise
                active = False  # it stopped meanwhile
            else:
                self.stream = stream  # first: the callback may end it at once
                stream.eventAddCallback(libvirt.VIR_STREAM_EVENT_READABLE, from_vm, self)
        if not active and self.active:
            self.close_stream()
            say("The VM is not running: its console opens again when it starts")
        self.active = active

    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 <= 0:  # -2: busy; 0 cannot advance either
                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 main():
    # what the operator reads goes to socketiod's log (the journal, docker
    # logs), on the descriptor it names; the user's terminal gets a short line
    if os.environ.get("CONSOLECALLBACK_LOG_FD"):
        log.addHandler(logging.StreamHandler(os.fdopen(int(os.environ["CONSOLECALLBACK_LOG_FD"]), "w", buffering=1)))
        log.propagate = False
    instance = Instance.objects.select_related("compute").filter(pk=int(sys.argv[1])).first()
    if instance is None:  # deleted since socketiod accepted the token
        say("The VM is no longer in the panel")
        sys.exit(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)
        console.check()
    except libvirt.libvirtError as err:
        log.warning("serial console of VM %s did not open: %s", instance.name, 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):
                    log.info("serial console of VM %s closed: its user may no longer open it", instance.name)
                    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
        log.warning("serial console of VM %s closed: %s", instance.name, err)
        say("The console of %s closed: %s" % (instance.name, public_error_text(err)))
        sys.exit(1)


if __name__ == "__main__":
    main()