From d40a20138b3f44c635bd738ba4da377c4ed9c5e8 Mon Sep 17 00:00:00 2001
From: Reid 'arrdem' McKenzie <me@arrdem.com>
Date: Mon, 29 May 2023 10:27:52 -0600
Subject: [PATCH] Get PrusaSlicer working; fix some auth problems

---
 .../src/python/tentacles/__main__.py          | 16 ++++-
 .../src/python/tentacles/blueprints/api.py    | 69 +++++++++++++++----
 .../python/tentacles/blueprints/file_ui.py    | 12 ++--
 .../src/python/tentacles/blueprints/job_ui.py |  6 +-
 .../python/tentacles/blueprints/printer_ui.py |  6 +-
 .../src/python/tentacles/blueprints/util.py   |  6 +-
 6 files changed, 86 insertions(+), 29 deletions(-)

diff --git a/projects/tentacles/src/python/tentacles/__main__.py b/projects/tentacles/src/python/tentacles/__main__.py
index 80d66d7..e61ad32 100644
--- a/projects/tentacles/src/python/tentacles/__main__.py
+++ b/projects/tentacles/src/python/tentacles/__main__.py
@@ -51,8 +51,20 @@ def create_j2_request_global(app):
 
 
 def user_session():
-    if (session_id := request.cookies.get("sid", "")) and (
-        uid := ctx.db.try_key(session_id)
+    if (
+        (
+            (session_id := request.cookies.get("sid", ""))
+            and (uid := ctx.db.try_key(session_id))
+        )
+        or (
+            request.authorization
+            and request.authorization.token
+            and (uid := ctx.db.try_key(request.authorization.token))
+        )
+        or (
+            (api_key := request.headers.get("x-api-key"))
+            and (uid := ctx.db.try_key(api_key))
+        )
     ):
         ctx.sid = session_id
         ctx.uid = uid
diff --git a/projects/tentacles/src/python/tentacles/blueprints/api.py b/projects/tentacles/src/python/tentacles/blueprints/api.py
index 53f3007..5d312db 100644
--- a/projects/tentacles/src/python/tentacles/blueprints/api.py
+++ b/projects/tentacles/src/python/tentacles/blueprints/api.py
@@ -7,29 +7,47 @@ import os
 from typing import Optional
 
 from flask import Blueprint, current_app, request
-from tentacles.blueprints.util import (
-    requires_admin,
-    requires_auth,
-)
 from tentacles.globals import ctx
 
 
 BLUEPRINT = Blueprint("api", __name__, url_prefix="/api")
 
 
+def requires_admin(f):
+    def _helper(*args, **kwargs):
+        if not ctx.is_admin:
+            return {"error": "unauthorized"}, 401
+        else:
+            return f(*args, **kwargs)
+
+    _helper.__name__ = f.__name__
+    return _helper
+
+
+def requires_auth(f):
+    def _helper(*args, **kwargs):
+        if not ctx.uid:
+            return {"error": "unauthorized"}, 401
+        else:
+            return f(*args, **kwargs)
+
+    _helper.__name__ = f.__name__
+    return _helper
+
+
 ####################################################################################################
 # Printers
 #
 # The trick here is handling multipart uploads.
-@requires_admin
 @BLUEPRINT.route("/printers", methods=["POST"])
+@requires_admin
 def create_printer():
     pass
 
 
-@requires_auth
 @BLUEPRINT.route("/printers", methods=["GET"])
 @BLUEPRINT.route("/printers/", methods=["GET"])
+@requires_auth
 def list_printers():
     return {
         "printers": [
@@ -38,8 +56,8 @@ def list_printers():
     }, 200
 
 
-@requires_admin
 @BLUEPRINT.route("/printers", methods=["DELETE"])
+@requires_admin
 def delete_printer():
     pass
 
@@ -48,9 +66,9 @@ def delete_printer():
 # Files
 #
 # The trick here is handling multipart uploads.
-@requires_auth
 @BLUEPRINT.route("/files", methods=["POST"])
 @BLUEPRINT.route("/files/<location>", methods=["POST"])
+@requires_auth
 def create_file(location: Optional[str] = None):
     # This is the important one, because it's the one that PrusaSlicer et all use to upload jobs.
 
@@ -68,20 +86,29 @@ def create_file(location: Optional[str] = None):
 
     else:
         digest = sha3_256()
+        digest.update(
+            f"$USER${ctx.uid}$".encode()
+        )  # Salt filenames per-user to avoid collisions
         digest.update(file.filename.encode())
         sanitized_filename = digest.hexdigest() + ".gcode"
         sanitized_path = os.path.join(
             current_app.config["UPLOAD_FOLDER"], sanitized_filename
         )
+        if os.path.exists(sanitized_path):
+            return {"error": "file exists already"}, 409
+
         file.save(sanitized_path)
-        ctx.db.create_file(ctx.uid, file.filename, sanitized_path)
+        fid = ctx.db.create_file(ctx.uid, file.filename, sanitized_path)
+
+        if request.form.get("print", "").lower() == "true":
+            ctx.db.create_job(ctx.uid, fid)
 
         return {"status": "ok"}, 202
 
 
-@requires_auth
 @BLUEPRINT.route("/files", methods=["GET"])
 @BLUEPRINT.route("/files/", methods=["GET"])
+@requires_auth
 def get_files():
     return {
         "files": [
@@ -97,22 +124,22 @@ def get_files():
     }, 200
 
 
-@requires_auth
 @BLUEPRINT.route("/files", methods=["DELETE"])
+@requires_auth
 def delete_file():
     pass
 
 
 ####################################################################################################
 # Jobs
-@requires_auth
 @BLUEPRINT.route("/jobs", methods=["POST"])
+@requires_auth
 def create_job():
     pass
 
 
-@requires_auth
 @BLUEPRINT.route("/jobs", methods=["GET"])
+@requires_auth
 def get_jobs():
     return {
         "jobs": [
@@ -128,15 +155,27 @@ def get_jobs():
     }, 200
 
 
-@requires_auth
 @BLUEPRINT.route("/jobs", methods=["DELETE"])
+@requires_auth
 def delete_job():
     pass
 
 
 ####################################################################################################
 # API tokens
-@requires_auth
 @BLUEPRINT.route("/tokens", methods=["GET"])
+@requires_auth
 def get_tokens():
     pass
+
+
+@BLUEPRINT.route("/version", methods=["GET"])
+@requires_auth
+def get_version():
+    """OctoPrint compatability endpoint."""
+
+    return {
+        "api": "0.1",
+        "server": "1.3.10",
+        "text": "OctoPrint 1.3.10 (Tentacles)",
+    }, 200
diff --git a/projects/tentacles/src/python/tentacles/blueprints/file_ui.py b/projects/tentacles/src/python/tentacles/blueprints/file_ui.py
index f98c42e..4555b09 100644
--- a/projects/tentacles/src/python/tentacles/blueprints/file_ui.py
+++ b/projects/tentacles/src/python/tentacles/blueprints/file_ui.py
@@ -20,8 +20,8 @@ log = logging.getLogger(__name__)
 BLUEPRINT = Blueprint("files", __name__)
 
 
-@requires_auth
 @BLUEPRINT.route("/files", methods=["GET"])
+@requires_auth
 def list_files():
     if request.method == "POST":
         flash("Not supported yet", category="warning")
@@ -29,8 +29,8 @@ def list_files():
     return render_template("files.html.j2")
 
 
-@requires_auth
 @BLUEPRINT.route("/files", methods=["POST"])
+@requires_auth
 def manipulate_files():
     match request.form.get("action"):
         case "upload":
@@ -43,11 +43,15 @@ def manipulate_files():
 
         case "delete":
             file = ctx.db.fetch_file(ctx.uid, int(request.form.get("file_id")))
-            if any(job.finished_at is None for job in ctx.db.list_jobs_by_file(file.id)):
+            if any(
+                job.finished_at is None for job in ctx.db.list_jobs_by_file(file.id)
+            ):
                 flash("File is in use", category="error")
                 return render_template("files.html.j2"), 400
 
-            os.unlink(file.path)
+            if os.path.exists(file.path):
+                os.unlink(file.path)
+
             ctx.db.delete_file(ctx.uid, file.id)
             flash("File deleted", category="info")
 
diff --git a/projects/tentacles/src/python/tentacles/blueprints/job_ui.py b/projects/tentacles/src/python/tentacles/blueprints/job_ui.py
index 7cd9f85..ead52da 100644
--- a/projects/tentacles/src/python/tentacles/blueprints/job_ui.py
+++ b/projects/tentacles/src/python/tentacles/blueprints/job_ui.py
@@ -18,21 +18,21 @@ log = logging.getLogger(__name__)
 BLUEPRINT = Blueprint("jobs", __name__)
 
 
-@requires_auth
 @BLUEPRINT.route("/jobs", methods=["GET"])
+@requires_auth
 def list_jobs():
     return render_template("jobs.html.j2")
 
 
-@requires_auth
 @BLUEPRINT.route("/jobs", methods=["POST"])
+@requires_auth
 def manipulate_jobs():
     match request.form.get("action"):
         case "enqueue":
             ctx.db.create_job(ctx.uid, int(request.form.get("file_id")))
             flash("Job created!", category="info")
 
-        case  "duplicate":
+        case "duplicate":
             if job := ctx.db.fetch_job(ctx.uid, int(request.form.get("job_id"))):
                 ctx.db.create_job(ctx.uid, job.file_id)
                 flash("Job created!", category="info")
diff --git a/projects/tentacles/src/python/tentacles/blueprints/printer_ui.py b/projects/tentacles/src/python/tentacles/blueprints/printer_ui.py
index c9f9d29..3c71ebc 100644
--- a/projects/tentacles/src/python/tentacles/blueprints/printer_ui.py
+++ b/projects/tentacles/src/python/tentacles/blueprints/printer_ui.py
@@ -18,14 +18,14 @@ log = logging.getLogger(__name__)
 BLUEPRINT = Blueprint("printer", __name__)
 
 
-@requires_admin
 @BLUEPRINT.route("/printers")
+@requires_admin
 def printers():
     return render_template("printers.html.j2")
 
 
-@requires_admin
 @BLUEPRINT.route("/printers/add", methods=["get", "post"])
+@requires_admin
 def add_printer():
     if not is_logged_in():
         return redirect("/")
@@ -50,7 +50,7 @@ def add_printer():
     return render_template("add_printer.html.j2")
 
 
-@requires_admin
 @BLUEPRINT.route("/printers/delete")
+@requires_admin
 def delete_printer():
     return render_template("delete_printer.html.j2")
diff --git a/projects/tentacles/src/python/tentacles/blueprints/util.py b/projects/tentacles/src/python/tentacles/blueprints/util.py
index be30237..efeb5d1 100644
--- a/projects/tentacles/src/python/tentacles/blueprints/util.py
+++ b/projects/tentacles/src/python/tentacles/blueprints/util.py
@@ -21,10 +21,11 @@ def requires_admin(f):
     def _helper(*args, **kwargs):
         if not ctx.is_admin:
             flash("Sorry, admins only", category="error")
-            redirect("/")
+            return redirect("/")
         else:
             return f(*args, **kwargs)
 
+    _helper.__name__ = f.__name__
     return _helper
 
 
@@ -32,8 +33,9 @@ def requires_auth(f):
     def _helper(*args, **kwargs):
         if not ctx.uid:
             flash("Please log in first", category="error")
-            redirect("/")
+            return redirect("/")
         else:
             return f(*args, **kwargs)
 
+    _helper.__name__ = f.__name__
     return _helper