From cba94e8b309e7d7f2383b56aa9a0f10e08ca0615 Mon Sep 17 00:00:00 2001
From: White <white@whites.macbook.pro.hungrysam>
Date: Sat, 7 Dec 2024 02:00:37 -0800
Subject: [PATCH] reworked code to address path traversal vuln

---
 PVEDiscordDark/serve.py | 67 ++++++++++++++++++-----------------------
 1 file changed, 30 insertions(+), 37 deletions(-)

diff --git a/PVEDiscordDark/serve.py b/PVEDiscordDark/serve.py
index 84734c8..9c90644 100644
--- a/PVEDiscordDark/serve.py
+++ b/PVEDiscordDark/serve.py
@@ -1,63 +1,56 @@
-# Script to assist with PVEDiscordDark development
-#
-# By default serves HTTP on port 3000, any *.js request gets the JS script, any *.css request gets the CSS file and any image request gets corresponding image
-# Meant to be used with the "Requestly" browser extension to redirect PVEDD requests from PVE server to localhost:3000
-#
-
 from http.server import HTTPServer, BaseHTTPRequestHandler
 import json
 import os
+import mimetypes
 
 PORT = 3000
 DIR_SASS = os.path.join(os.path.dirname(__file__), "sass")
 DIR_IMAGES = os.path.join(os.path.dirname(__file__), "images")
 DIR_JS = os.path.join(os.path.dirname(__file__), "js")
 
+STATUS_OK = 200
+STATUS_NOT_FOUND = 404
+STATUS_BAD_REQUEST = 400
 
 class Server(BaseHTTPRequestHandler):
     def log_message(self, format, *args):
-        return
-
-    def _set_headers(self, status, type):
+        # Basic logging to console
+        print(format % args)
+        
+    def _send_response(self, status, content_type):
         self.send_response(status)
-        self.send_header("Content-type", type)
+        self.send_header("Content-type", content_type)
         self.end_headers()
-
+        
+    def _read_file(self, directory, filename, default_status, content_type):
+        try:
+            with open(os.path.join(directory, filename), "rb") as f:
+                return f.read(), STATUS_OK, content_type
+        except FileNotFoundError:
+            return None, STATUS_NOT_FOUND, "application/json"
+        
     def do_GET(self):
-        status = 200
-        type = "application/json"
-        data = None
-
         file = self.path.rpartition("/")[2]
         ext = file.rpartition(".")[2]
-
+        
         if ext == "css":
-            data = open(os.path.join(DIR_SASS, "PVEDiscordDark.css"), "rb").read()
-            type = "text/css"
+            data, status, content_type = self._read_file(DIR_SASS, "PVEDiscordDark.css", STATUS_NOT_FOUND, "text/css")
         elif ext == "js":
-            data = open(os.path.join(DIR_JS, "PVEDiscordDark.js"), "rb").read()
-            type = "application/javascript"
-        elif ext == "png" or ext == "jpg" or ext == "jpeg":
-            try:
-                data = open(os.path.join(DIR_IMAGES, file), "rb").read()
-                type = f"image/{ext}"
-            except FileNotFoundError:
-                status = 404
-        elif ext == "svg":
-            try:
-                data = open(os.path.join(DIR_IMAGES, file), "rb").read()
-                type = f"image/svg+xml"
-            except FileNotFoundError:
-                status = 404
+            data, status, content_type = self._read_file(DIR_JS, "PVEDiscordDark.js", STATUS_NOT_FOUND, "application/javascript")
+        elif ext in ["png", "jpg", "jpeg", "svg"]:
+            content_type, _ = mimetypes.guess_type(file)
+            data, status, content_type = self._read_file(DIR_IMAGES, file, STATUS_NOT_FOUND, content_type)
         else:
-            status = 400
-        self._set_headers(status, type)
-        if status == 200:
-            self.wfile.write(data)
+            data, status, content_type = None, STATUS_BAD_REQUEST, "application/json"
+        
+        self._send_response(status, content_type)
+        
+        if status == STATUS_OK:
+            if data is not None:
+                self.wfile.write(data)
         else:
             self.wfile.write(json.dumps({"error": status}).encode())
 
-
 if __name__ == "__main__":
     print(f"Serving on localhost:{PORT}")
     server = HTTPServer(server_address=("", PORT), RequestHandlerClass=Server)