Add system for handling jobs (like file upload/download etc.)

1 files changed, 38 insertions, 0 deletions

diff --git a/session.py b/session.py
index 40a091f..5e4e3f9 100644
--- a/session.py
+++ b/session.py
@@ -3,6 +3,7 @@ import random
 import secrets
 import signal
 import string
+import subprocess
 import threading
 
 import psutil
@@ -18,12 +19,23 @@ 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 ServerThread(threading.Thread):
     def __init__(self, app, port: int):
         super().__init__()
@@ -46,7 +58,9 @@ class Session:
     STATUS_JOINED = "active"
 
     def __init__(self, mockup_session: bool):
+        self.done_jobs = []
         self.mockup_session = mockup_session
+        self.desktop_dir = get_desktop_dir()
         self._generate_password()
         self._start_vnc()
         self._start_trigger_service()
@@ -158,6 +172,30 @@ class Session:
                 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."""