path: root/session.py

import os
import random
import secrets
import signal
import string
import threading

import psutil
import requests
from werkzeug.serving import make_server

import port_for
from flask import Flask, abort, request
from vnc import run_vnc, save_password

API_SERVER = "http://127.0.0.1:8000"
BASE_URL = API_SERVER + "/app/rwa/api/"
REGISTER_URL = BASE_URL + "register/"
STOP_URL = BASE_URL + "stop/"
STATUS_URL = BASE_URL + "status/"


def random_digits(length: int):
    return "".join(random.choice(string.digits) for _ in range(length))


class ServerThread(threading.Thread):
    def __init__(self, app, port: int):
        super().__init__()
        self.srv = make_server("127.0.0.1", port, app)
        self.ctx = app.app_context()
        self.ctx.push()

    def run(self):
        self.srv.serve_forever()

    def shutdown(self):
        self.srv.shutdown()


class Session:
    #: Session is running
    STATUS_RUNNING = "running"

    #: Remote has joined the session
    STATUS_JOINED = "active"

    def __init__(self, mockup_session: bool):
        self.mockup_session = mockup_session
        self._generate_password()
        self._start_vnc()
        self._start_trigger_service()
        self._register_session()
        self.status_text = self.STATUS_RUNNING

    @property
    def pid(self) -> int:
        return self.vnc_pid

    @property
    def port(self) -> int:
        return self.ws_port

    @property
    def _api_headers(self) -> dict:
        return {"Authorization": f"Token {self.api_token}"}

    def _generate_password(self):
        """Generate password for x11vnc and save it."""
        self.password = secrets.token_urlsafe(20)

        # Don't actually save a password if we just pretend to be a session.
        if not self.mockup_session:
            self.pw_filename = save_password(self.password)

    def _start_vnc(self):
        """Start x11vnc server if not in mockup_session mode."""
        if not self.mockup_session:
            process_info = run_vnc(self.pw_filename)

            self.vnc_pid = process_info["vnc"]["pid"]
            self.vnc_port = process_info["vnc"]["port"]
            self.ws_pid = process_info["ws"]["pid"]
            self.ws_port = process_info["ws"]["port"]
        else:
            self.ws_port = port_for.select_random()
            self.vnc_port = port_for.select_random()

            # Use negative values to ensure we don't do something harmful
            # to random processes
            self.ws_pid = int("-" + random_digits(5))
            self.vnc_pid = int("-" + random_digits(5))

            # Create a temporary file to indicate that this process is still 'Running'
            filename = f"/tmp/rwa/{str(self.ws_port) + str(self.vnc_port) + str(self.ws_pid) + str(self.vnc_pid)}.lock"
            new_file = open(filename, "w")
            new_file.write("this session is running")

    def _register_session(self):
        """Register session in RWA if not in mockup_session mode."""
        if not self.mockup_session:
            r = requests.post(
                REGISTER_URL,
                json={
                    "port": self.ws_port,
                    "password": self.password,
                    "pid": self.vnc_pid,
                    "trigger_port": self.trigger_port,
                    "trigger_token": self.trigger_token,
                },
            )
            print(r)
            self.meta = r.json()
            self.session_id = self.meta["session_id"]
            self.web_url = self.meta["url"]
            self.api_token = self.meta["token"]
            self.pin = self.meta["pin"]
        else:
            print('"Registered" in RWA')
            self.meta = {}
            self.session_id = int(random_digits(10))
            self.web_url = "testhostname:" + random_digits(5) + "/RWA/test/"
            self.api_token = secrets.token_urlsafe(10)
            self.pin = int(random_digits(5))

    def _start_trigger_service(self):
        self.trigger_port = port_for.select_random()
        self.trigger_token = secrets.token_urlsafe(20)

        app = Flask(__name__)

        @app.route("/", methods=["POST"])
        def trigger():
            json = request.json
            if json.get("token", "") == self.trigger_token:
                self._trigger()
                return "Successful triggered"
            else:
                return abort(403)

        self.trigger_thread = ServerThread(app=app, port=self.trigger_port)
        self.trigger_thread.start()

    def _trigger(self):
        """Event triggered by Django."""
        print("Triggered")
        self.pull()

    def pull(self):
        """Update status: Get status from Django."""
        if not self.mockup_session:
            r = requests.get(
                STATUS_URL, params={"id": self.session_id}, headers=self._api_headers
            )

            if r.status_code in (401, 402, 403, 404, 405):
                # Session doesn't exist anymore, so stop it local
                self.stop(triggered=True)
            else:
                self.status_text = r.json()["status"]

    def push(self):
        """Update status: Push status to Django."""
        pass

    def stop(self, triggered: bool = False):
        """Stop session and clean up."""
        if self.mockup_session:
            filename = f"/tmp/rwa/{str(self.ws_port) + str(self.vnc_port) + str(self.ws_pid) + str(self.vnc_pid)}.lock"
            if os.path.isfile(filename):
                os.remove(filename)

            # Delete self
            del self
            return

        # Kill VNC
        if self.vnc_pid in psutil.pids():
            print("Kill VNC.")
            os.kill(self.vnc_pid, signal.SIGTERM)

        # Kill websockify
        if self.ws_pid in psutil.pids():
            print("Kill websockify.")
            os.kill(self.ws_pid, signal.SIGTERM)

        # Delete PW file
        if os.path.exists(self.pw_filename):
            print("Delete password file")
            os.remove(self.pw_filename)

        if hasattr(self, "trigger_thread"):
            print("Kill trigger service.")
            self.trigger_thread.shutdown()

        self.push()

        if not triggered:
            requests.post(
                STOP_URL, params={"id": self.session_id}, headers=self._api_headers
            )

        self.status_text = "stopped"

        # Delete self
        del self

    @property
    def vnc_process_running(self):
        """Check if the VNC process is still running."""
        if self.mockup_session:
            filename = f"/tmp/rwa/{str(self.ws_port) + str(self.vnc_port) + str(self.ws_pid) + str(self.vnc_pid)}.lock"
            return os.path.isfile(filename)

        if self.vnc_pid in psutil.pids():
            p = psutil.Process(self.vnc_pid)
            if p.status() == "zombie":
                return False
            return True
        return False

    @property
    def client_meta(self):
        return {"id": self.pid, "url": self.web_url, "pin": self.pin}

    @property
    def status(self):
        return {"id": self.pid, "status": self.status_text}