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 ) 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 ) # 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}