Prevent crashing when asking for status with invalid pid & mockup-mode: make sessionid's positive

2 files changed, 19 insertions, 8 deletions

diff --git a/service.py b/service.py
index 1faae4d..4c4e927 100755
--- a/service.py
+++ b/service.py
@@ -47,14 +47,21 @@ class RWAService(dbus.service.Object):
         self._update_session(pid)
         return self._get_status(pid)
 
-    @dbus.service.method("org.ArcticaProject.RWA", in_signature="i")
+    @dbus.service.method("org.ArcticaProject.RWA", in_signature="i", out_signature="s")
     def stop(self, pid: int):
         """Stop a remote session."""
-        session = self.sessions[pid]
+        try:
+            session = self.sessions[pid]
+        except KeyError:
+            return json.dumps({"pid": pid, "status": "stopped"}, sort_keys=True)
         session.stop()
+        return json.dumps({"id": pid, "status": "stopped"}, sort_keys=True)
 
     def _get_status(self, pid: int) -> str:
-        session = self.sessions[pid]
+        try:
+            session = self.sessions[pid]
+        except KeyError:
+            return json.dumps({"id": pid, "status": "dead"}, sort_keys=True)
         return json.dumps(session.status)
 
     def _ensure_update_service(self):
@@ -65,7 +72,13 @@ class RWAService(dbus.service.Object):
 
     def _update_session(self, pid: int):
         """Update the status of a session."""
-        session = self.sessions[pid]
+        try:
+            session = self.sessions[pid]
+        except KeyError:
+            print(f"Session #{pid}")
+            print("Session is dead.")
+            return
+
         print(f"Session #{session.pid}")
 
         # Check if VNC process is still running
diff --git a/session.py b/session.py
index 40a091f..2f3accb 100644
--- a/session.py
+++ b/session.py
@@ -86,10 +86,8 @@ class Session:
             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))
+            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"