Simplify generation of random numbers

1 files changed, 10 insertions, 21 deletions

diff --git a/session.py b/session.py
index 6b248b8..9fc755d 100644
--- a/session.py
+++ b/session.py
@@ -18,6 +18,10 @@ REGISTER_URL = BASE_URL + "register/"
 STOP_URL = BASE_URL + "stop/"
 
 
+def random_digits(length: int):
+    return "".join(random.choice(string.digits) for _ in range(length))
+
+
 class Session:
     #: Session is running
     STATUS_RUNNING = "running"
@@ -67,12 +71,8 @@ class Session:
 
             # 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))
-            )
+            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"
@@ -101,21 +101,10 @@ class Session:
         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)))
+            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()