webrun

#!/usr/bin/python3
import asyncio
import argparse
import logging
import signal
import os
import sys
import shutil
import shlex
import json
import socket
from unitd import SimpleProcess
from unitd.processpool import ProcessPool
from unitd.signals import create_future_for_signal
import unitd.config
from unitd.config import HostPort
from unitd.process import create_future


XSERVERS = ("Xtigervnc", "Xvnc4")

log = logging.getLogger()


class Fail(RuntimeError):
    pass


def helper_config(config):
    res = unitd.config.Config()
    for skey in "working_directory", "user", "group":
        setattr(res.service, skey, getattr(config.service, skey))
    return res


class Xserver(SimpleProcess):
    def __init__(self, config, cmdline, loop=None):
        x_config = helper_config(config)
        x_config.service.syslog_identifier = "xserver"
        x_config.service.user = os.getuid()
        x_config.service.user = os.getgid()
        x_config.service.exec_start.append(cmdline)
        super().__init__(x_config, loop=loop)
        self.server_ready = None

    def _preexec(self):
        super()._preexec()
        signal.signal(signal.SIGUSR1, signal.SIG_IGN)
        #signal.signal(signal.SIGINT, signal.SIG_IGN)

    @asyncio.coroutine
    def _start(self):
        # When starting an X server, if the X server sees that SIGUSR1 is
        # ignored, it sends SIGUSR1 to the parent process when it is ready to
        # accept connections. We wait until we receive SIGUSR1 before declaring
        # that the X server has been started.
        self.server_ready = create_future_for_signal(signal.SIGUSR1, loop=self.loop)
        return (yield from super()._start())

    @asyncio.coroutine
    def _confirm_start(self):
        yield from self.server_ready
        yield from super()._confirm_start()


class Websockify(SimpleProcess):
    def __init__(self, config, cmdline, loop=None):
        wr_config = helper_config(config)
        wr_config.service.syslog_identifier = "websocksify"
        wr_config.service.exec_start.append(cmdline)
        super().__init__(wr_config)
        self.listen_port = config.webrun.web_port

    @asyncio.coroutine
    def _confirm_start(self):
        try:
            yield from self._wait_for_listen(self.listen_port.host, self.listen_port.port)
            yield from super()._confirm_start()
        except asyncio.CancelledError:
            self._start_confirmed = False

    @asyncio.coroutine
    def _wait_for_listen(self, host, port):
        """
        Try to connect to a (host, port) until something starts listening on it
        """
        sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
        # In 3.5, sock_connect can take names to be resolved.
        # In 3.4, we need to resolve them first.
        resolved = yield from self.loop.getaddrinfo(host, port, family=socket.AF_INET, type=socket.SOCK_STREAM)
        try:
            while True:
                try:
                    yield from self.loop.sock_connect(sock, resolved[0][4])
                except ConnectionRefusedError:
                    yield from asyncio.sleep(0.2)
                else:
                    return
        finally:
            sock.close()


class RemoteApplication:
    """
    Run an application inside a private VNC server, and export the VNC
    connection to it over websocket
    """
    def __init__(self, config, desktop_title="Remote application"):
        self.config = config
        self.desktop_title = desktop_title

        for xserver in XSERVERS:
            self.xserver = xserver
            self.xserver_path = shutil.which(xserver)
            log.debug("X server lookup: tried %s, got %s", xserver, self.xserver_path)
            if self.xserver_path is not None:
                break
        log.debug("X server lookup: selected X server %s", self.xserver_path)
        if self.xserver_path is None:
            raise Fail("X server not found. Tried: " + ", ".join(XSERVERS))

        self.loop = asyncio.get_event_loop()
        self.xserver = self._create_xserver()
        self.xclient = self._create_xclient()
        self.websockify = self._create_websockify()
        self.started = create_future(self.loop)

    def _create_xserver(self):
        """
        Start the X server and return when it is ready to accept connections
        from X clients
        """
        cmd = [self.xserver_path, "-ac",
               "-geometry", self.config.webrun.display_geometry,
               "-desktop", str(self.config.webrun.display_number),
               "-SecurityTypes", "None",
               "-rfbport", str(self.config.webrun.vnc_port.port)]
        if self.xserver == "Xtigervnc":
            cmd += [
               "-interface", self.config.webrun.vnc_port.host,
            ]
        cmd += [
               ":{}".format(self.config.webrun.display_number)
        ]
        return Xserver(self.config, cmd, loop=self.loop)

    def _create_xclient(self):
        """
        Start the X client
        """
        env = dict(os.environ)
        env["DISPLAY"] = ":{}".format(self.config.webrun.display_number)
        if self.config.service.syslog_identifier is None:
            self.config.service.syslog_identifier = "xclient"
        return SimpleProcess(self.config, env=env, loop=self.loop)

    def _create_websockify(self):
        """
        Start websockify to export the X client connection via a websocket
        server
        """
        cmd = ["websockify",
               "--run-once",
               "--web", "/usr/share/novnc/"]
        if self.config.webrun.web_connect_timeout:
            cmd += ["--timeout", str(self.config.webrun.web_connect_timeout)]
        cmd.append(str(self.config.webrun.web_port))
        cmd.append(str(self.config.webrun.vnc_port))
        return Websockify(self.config, cmd, loop=self.loop)

    @asyncio.coroutine
    def run(self):
        """
        Wait until any one of the X server, X application or websocket proxy
        quits, then make sure all the others quit as well
        """
        pool = ProcessPool(loop=self.loop)
        pool.set_quit_signal(signal.SIGINT)
        yield from pool.start_sync(self.xserver)
        yield from pool.start_sync(self.xclient)
        yield from pool.start_sync(self.websockify)
        self.started.set_result(pool.success)
        yield from pool.run()


def start_everything(config, args):
    runner = RemoteApplication(
            config=config,
            desktop_title=args.title)

    loop = asyncio.get_event_loop()
    try:
        result = {}
        try:
            log.debug("Starting processes")
            task = loop.create_task(runner.run())

            result["started"] = loop.run_until_complete(runner.started)

            if result["started"]:
                log.info("Ready for connections, interrupt with ^C")

                result["url"] = "http://{addr}/vnc.html?host={addr.host}&port={addr.port}&autoconnect=true".format(addr=config.webrun.web_port)
                log.info("Connect using %s", result["url"])
                if args.open and not args.daemon:
                    with config.service.nonroot():
                        import webbrowser
                        webbrowser.open(result["url"])
        except Exception as e:
            result["started"] = False
            result["error"] = str(e)
        finally:
            print(json.dumps(result))
            if args.daemon:
                sys.stdout.close()

        if args.daemon:
            devnull_fd = os.open("/dev/null", os.O_RDWR)
            os.dup2(devnull_fd, 0)
            os.dup2(devnull_fd, 1)
            os.dup2(devnull_fd, 2)
            os.umask(0o27)

    finally:
        loop.run_until_complete(task)
        loop.close()


def main():
    config = unitd.config.Config()

    parser = argparse.ArgumentParser(
            description="start X client on a private X server exported via VNC on the web")
    parser.add_argument("--verbose", "-v", action="store_true", help="verbose output")
    parser.add_argument("--debug", action="store_true", help="debug output")
    parser.add_argument("--logfile", action="store", help="lot to the given file")
    parser.add_argument("--geometry", action="store", help="X display geometry")
    parser.add_argument("--display", action="store", type=int, default=None, help="X display number, default: {}".format(config.webrun.display_number))
    parser.add_argument("--title", action="store", default="Desktop session", help="desktop session title")
    parser.add_argument("--vnc-port", action="store", default=None, help="vnc server port, default: 5900 + display number")
    parser.add_argument("--web-port", action="store", default=None, help="websocket server port, default: 6080")
    parser.add_argument("--open", action="store_true", help="open application in browser")
    parser.add_argument("--daemon", action="store_true", help="daemonize")
    parser.add_argument("--service", action="store_true", help="argument is a .service file rather than a command line")
    parser.add_argument("--web-connect-timeout", action="store", type=int, default=None, help="quit if there is no web connection after this amount of seconds, default: no timeout")
    parser.add_argument("cmd", nargs="+", help="command and arguments to run")

    args = parser.parse_args()

    log_format = "%(asctime)-15s %(levelname)s %(message)s"
    level = logging.WARN
    if args.debug:
        level = logging.DEBUG
    elif args.verbose:
        level = logging.INFO

    if args.logfile:
        config.webrun.log_file = args.logfile

    if args.service:
        if len(args.cmd) != 1:
            raise Fail("When using --service only provide one argument: the pathname to the .service file")
        config.read_file(args.cmd[0])
    else:
        config.service.exec_start.append(args.cmd)

    if config.webrun.log_file:
        logging.basicConfig(level=level, filename=config.webrun.log_file, format=log_format)
    else:
        logging.basicConfig(level=level, stream=sys.stderr, format=log_format)

    if args.display is not None:
        config.webrun.display_number = args.display
    if args.geometry is not None:
        config.webrun.display_geometry = args.geometry
    if args.vnc_port is not None:
        config.webrun.vnc_port = HostPort.parse(args.vnc_port)
    if args.web_port is not None:
        config.webrun.web_port = HostPort.parse(args.web_port)
    if args.web_connect_timeout is not None:
        config.webrun.web_connect_timeout = args.web_connect_timeout

    if args.daemon:
        fd_r, fd_w = os.pipe()
        pid = os.fork()
        if pid != 0:
            with config.service.nonroot():
                os.dup2(fd_r, 0)
                os.close(fd_w)
                # Parent process
                # Wait for child
                child_info = sys.stdin.read()
                sys.stdout.write(child_info)
                if args.open:
                    with config.service.nonroot():
                        result = json.loads(child_info)
                        import webbrowser
                        webbrowser.open(result["url"])
                return
        os.dup2(fd_w, 1)
        os.close(fd_w)

    if args.daemon:
        try:
            os.setsid()
        except PermissionError:
            pass
        os.chdir("/")

    start_everything(config, args)


if __name__ == "__main__":
    try:
        main()
    except Fail as e:
        print(e, file=sys.stderr)
        sys.exit(1)
    except Exception:
        log.exception("uncaught exception")