diff options
Diffstat (limited to 'session.py')
| -rw-r--r-- | session.py | 92 |
1 files changed, 67 insertions, 25 deletions
@@ -3,6 +3,8 @@ import secrets import signal from multiprocessing import Process +import random, string + import psutil import requests @@ -23,7 +25,8 @@ class Session: #: Remote has joined the session STATUS_JOINED = "joined" - def __init__(self): + def __init__(self, mockup_session: bool): + self.mockup_session = mockup_session self._generate_password() self._start_vnc() self._start_trigger_service() @@ -44,35 +47,61 @@ class Session: def _generate_password(self): """Generate password for x11vnc and save it.""" self.password = secrets.token_urlsafe(20) - self.pw_filename = save_password(self.password) + + # 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.""" - process_info = run_vnc(self.pw_filename) + """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('-' + ''.join(random.choice(string.digits) for _ in range(5))) + self.vnc_pid = int('-' + ''.join(random.choice(string.digits) for _ in range(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') - 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"] def _register_session(self): - """Register session in RWA.""" - 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"] + """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(''.join(random.choice(string.digits) for _ in range(10))) + self.web_url = "testhostname:" + ''.join(random.choice(string.digits) for _ in range(5)) + "/RWA/test/" + self.api_token = ''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(10)) + self.pin = int(''.join(random.choice(string.digits) for _ in range(5))) def _start_trigger_service(self): self.trigger_port = port_for.select_random() @@ -109,6 +138,15 @@ class Session: 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.") @@ -141,6 +179,10 @@ class Session: @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": |