Generate PIN in this service and provide authentication mechanisms

3 files changed, 79 insertions, 45 deletions

diff --git a/service.py b/service.py
index d77ac7c..ada9b15 100755
--- a/service.py
+++ b/service.py
@@ -33,7 +33,7 @@ import signal
 import tempfile
 import time
 from threading import Thread
-from typing import Union
+from typing import Union, Any
 
 import dbus
 import dbus.mainloop.glib
@@ -238,17 +238,17 @@ class RWASupportSessionService(dbus.service.Object):
         if self.one_time:
             self._stop_all()
 
-    def _trigger(self, token: str) -> bool:
+    def _trigger(self, session_id: int,  data: dict, method: str = "trigger") -> Union[dict, bool]:
         """Trigger a specific session via trigger token."""
-        logging.info(f"Triggered with token {token}")
+        logging.info(f"Triggered with session ID {session_id} and {data}")
 
         for session in self.sessions.values():
-            if token == session.trigger_token:
-                logging.info(f"Session #{session.pid} matches the token.")
-                session.trigger()
-                return True
+            if session.session_id == session_id:
+                r = session.trigger(data, method)
+                logging.info(f"Session #{session.pid} matches the ID: {r}")
+                return r
 
-        logging.warning("  No matching session found for this token.")
+        logging.warning("  No matching session found for this ID.")
         return False
 
     def _stop_all(self):
diff --git a/session.py b/session.py
index fe67fda..f977669 100644
--- a/session.py
+++ b/session.py
@@ -29,7 +29,7 @@ import secrets
 import signal
 import string
 import subprocess
-from typing import Dict, Union
+from typing import Dict, Union, Any
 
 import port_for
 import psutil
@@ -137,10 +137,8 @@ class Session:
                     REGISTER_URL,
                     json={
                         "port": self.ws_port,
-                        "password": self.password,
                         "pid": self.vnc_pid,
                         "trigger_port": self.trigger_port,
-                        "trigger_token": self.trigger_token,
                     },
                 )
             except requests.exceptions.ConnectionError:
@@ -157,29 +155,40 @@ class Session:
             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:
             logging.info(f"The session has pretended that he had created a session.")
             self.meta = {}
             self.session_id = int(random_digits(10))
-            self.web_url = "testhostname:" + random_digits(5) + "/RWA.Support/test/"
+            self.web_url = "http://example.com:" + random_digits(5) + "/app/rwasupport/test/"
             self.api_token = secrets.token_urlsafe(10)
-            self.pin = int(random_digits(5))
+        self.pin = int(random_digits(4))
 
-    def trigger(self):
+    def trigger(self, data: dict, method: str = "trigger") -> Union[dict, bool]:
         """Event triggered by Django."""
-        self.pull()
+        if method == "trigger" and data.get("token", "") == self.trigger_token:
+            self.pull()
+            return True
+        elif method == "authenticate" and data.get("pin", "") == self.pin:
+            return {
+                "password": self.password,
+                "trigger_token": self.trigger_token,
+            }
+
 
+        return False
     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
-            )
+            try:
+                r = requests.get(
+                    STATUS_URL, params={"id": self.session_id}, headers=self._api_headers
+                )
 
-            logging.info(
-                f"The session has received its status from RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
-            )
+                logging.info(
+                    f"The session has received its status from RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
+                )
+            except requests.ConnectionError:
+                pass
 
             if r.status_code in (401, 402, 403, 404, 405):
                 # Session doesn't exist anymore, so stop it local
@@ -208,14 +217,17 @@ class Session:
     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"])
-        r = requests.post(
-            MARK_JOB_AS_DONE_URL,
-            params={"id": job["job_id"]},
-            headers=self._api_headers,
-        )
-        logging.info(
-            f"The session has marked the job {job} as done in RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
-        )
+        try:
+            r = requests.post(
+                MARK_JOB_AS_DONE_URL,
+                params={"id": job["job_id"]},
+                headers=self._api_headers,
+            )
+            logging.info(
+                f"The session has marked the job {job} as done in RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
+            )
+        except requests.ConnectionError:
+            pass
 
     def push(self):
         """Update status: Push status to Django."""
@@ -251,12 +263,15 @@ class Session:
         self.push()
 
         if not triggered:
-            r = requests.post(
-                STOP_URL, params={"id": self.session_id}, headers=self._api_headers
-            )
-            logging.info(
-                f"The stop action has been registered in RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
-            )
+            try:
+                r = requests.post(
+                    STOP_URL, params={"id": self.session_id}, headers=self._api_headers
+                )
+                logging.info(
+                    f"The stop action has been registered in RWA.Support.WebApp with status code {r.status_code} and response {r.content.decode()}."
+                )
+            except requests.ConnectionError:
+                pass
 
         self.status_text = "stopped"
 
@@ -279,7 +294,7 @@ class Session:
 
     @property
     def client_meta(self) -> Dict[str, Union[str, int]]:
-        return {"id": self.pid, "url": self.web_url, "pin": self.pin}
+        return {"id": self.pid, "session_id": self.session_id, "url": self.web_url, "pin": self.pin}
 
     @property
     def status(self) -> Dict[str, Union[str, int]]:
diff --git a/trigger.py b/trigger.py
index 7d2b7e3..fa9734a 100644
--- a/trigger.py
+++ b/trigger.py
@@ -24,17 +24,17 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import threading
-from typing import Any, Callable
+from typing import Any, Callable, Optional, Union
 from wsgiref.simple_server import make_server
 
 import port_for
-from flask import Flask, abort, request
+from flask import Flask, abort, request, jsonify
 
 
 class TriggerServerThread(threading.Thread):
     """Simple Flask server (wrapped as thread) for triggering actions on sessions."""
 
-    def __init__(self, trigger_method: Callable[[str], Any]):
+    def __init__(self, trigger_method: Callable[[int, dict, Optional[str]], Union[dict, bool]]):
         super().__init__()
         self.port = port_for.select_random()
 
@@ -44,11 +44,30 @@ class TriggerServerThread(threading.Thread):
         def trigger():
             json = request.json
             token = json.get("token", "")
-            r = trigger_method(token)
-            if r:
-                return "Successful triggered"
-            else:
-                return abort(403)
+            try:
+                session_id = int(json.get("session_id"))
+                r = trigger_method(session_id, {"token": token}, "trigger")
+                if r:
+                    return "Successful triggered"
+                else:
+                    return abort(403)
+            except (ValueError, TypeError):
+                return abort(404)
+
+        @app.route("/authenticate/", methods=["POST"])
+        def authenticate():
+            json = request.json
+            try:
+                session_id = int(json.get("session_id"))
+                pin = int(json.get("pin", ""))
+                r = trigger_method(session_id, {"pin": pin}, "authenticate")
+                if r:
+                    return jsonify(r)
+                else:
+                    return abort(403)
+            except (ValueError, TypeError):
+                return abort(404)
+
 
         self.srv = make_server("0.0.0.0", self.port, app)
         self.ctx = app.app_context()