import os import random import secrets import signal import string import subprocess 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/" MARK_JOB_AS_DONE_URL = BASE_URL + "jobs/mark_as_done/" def random_digits(length: int): return "".join(random.choice(string.digits) for _ in range(length)) def get_desktop_dir(): """Get desktop directory from xdg vars.""" return ( subprocess.check_output(["xdg-user-dir", "DESKTOP"]) .decode() .strip() .replace("\n", "") ) class Session: #: Session is running STATUS_RUNNING = "running" #: Remote has joined the session STATUS_JOINED = "active" def __init__(self, trigger_port: int, mockup_session: bool = False): self.trigger_token = secrets.token_urlsafe(20) self.trigger_port = trigger_port self.done_jobs = [] self.mockup_session = mockup_session self.desktop_dir = get_desktop_dir() self.desktop_dir = get_desktop_dir() self._generate_password() self._start_vnc() 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() 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 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"] self.jobs = r.json()["jobs"] self._do_jobs() def _do_jobs(self): """Go through all jobs and execute undone ones.""" for job in self.jobs: if not job["done"] or job["job_id"] in self.done_jobs: job_type = job["job_type"] if job_type == "file": self._do_file_job(job) def _do_file_job(self, job): """Download a file from server to the user's desktop.""" subprocess.Popen(["wget", job["file"], "-P", self.desktop_dir]) self._mark_job_as_done(job) def _mark_job_as_done(self, job): """Mark a job as done (in this service and on the server).""" self.done_jobs.append(job["job_id"]) requests.post( MARK_JOB_AS_DONE_URL, params={"id": job["job_id"]}, headers=self._api_headers, ) 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}