#!/usr/bin/env python3
# consolecallback <instance 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 instances.models import Instance  # noqa: E402
from vrtManager.connection import connection_manager  # noqa: E402


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:
            self.stream = self.connection.newStream(libvirt.VIR_STREAM_NONBLOCK)
            self.domain.openConsole(None, self.stream, 0)  # a console in use is refused
            self.stream.eventAddCallback(libvirt.VIR_STREAM_EVENT_READABLE, from_vm, self)
        elif not active and self.stream is not None:
            self.close_stream()

    def close_stream(self):
        stream, self.stream = self.stream, None
        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
    elif console.stream is not None:
        try:
            console.stream.send(data)
        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():
    instance = Instance.objects.select_related("compute").get(pk=int(sys.argv[1]))
    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:
        sys.stdout.write("Cannot open the console of %s: %s\r\n" % (instance.name, err.get_error_message()))
        sys.exit(1)
    tty.setraw(0)
    libvirt.virEventAddHandle(0, libvirt.VIR_EVENT_HANDLE_READABLE, from_client, console)
    while console.running:
        time.sleep(0.2)
        console.check()


main()
