Fix: Sanitizing HTML inputs and Cors misconfiguration

Author: ARajgor

devika.py

@@ -26,7 +26,7 @@
 
 
 app = Flask(__name__)
-CORS(app)
+CORS(app, resources={r"/*": {"origins": ["https://localhost:3000"]}}) # Change the origin to your frontend URL
 app.register_blueprint(project_bp)
 socketio.init_app(app)
 
@@ -116,14 +116,6 @@ def get_agent_state():
     return jsonify({"state": agent_state})
 
 
-@app.route("/api/get-project-files/", methods=["GET"])
-@route_logger(logger)
-def project_files():
-    project_name = request.args.get("project_name")
-    files = AgentState.get_project_files(project_name)  
-    return jsonify({"files": files})
-
-
 @app.route("/api/get-browser-snapshot", methods=["GET"])
 @route_logger(logger)
 def browser_snapshot():

src/apis/project.py

@@ -1,4 +1,5 @@
 from flask import blueprints, request, jsonify, send_file, make_response
+from werkzeug.utils import secure_filename
 from src.logger import Logger, route_logger
 from src.config import Config
 from src.project import ProjectManager
@@ -13,20 +14,28 @@
 
 
 # Project APIs
+
+@project_bp.route("/api/get-project-files", methods=["GET"])
+@route_logger(logger)
+def project_files():
+    project_name = secure_filename(request.args.get("project_name"))
+    files = manager.get_project_files(project_name)  
+    return jsonify({"files": files})
+
 @project_bp.route("/api/create-project", methods=["POST"])
 @route_logger(logger)
 def create_project():
     data = request.json
     project_name = data.get("project_name")
-    manager.create_project(project_name)
+    manager.create_project(secure_filename(project_name))
     return jsonify({"message": "Project created"})
 
 
 @project_bp.route("/api/delete-project", methods=["POST"])
 @route_logger(logger)
 def delete_project():
     data = request.json
-    project_name = data.get("project_name")
+    project_name = secure_filename(data.get("project_name"))
     manager.delete_project(project_name)
     AgentState().delete_state(project_name)
     return jsonify({"message": "Project deleted"})
@@ -35,7 +44,7 @@ def delete_project():
 @project_bp.route("/api/download-project", methods=["GET"])
 @route_logger(logger)
 def download_project():
-    project_name = request.args.get("project_name")
+    project_name = secure_filename(request.args.get("project_name"))
     manager.project_to_zip(project_name)
     project_path = manager.get_zip_path(project_name)
     return send_file(project_path, as_attachment=False)
@@ -44,7 +53,7 @@ def download_project():
 @project_bp.route("/api/download-project-pdf", methods=["GET"])
 @route_logger(logger)
 def download_project_pdf():
-    project_name = request.args.get("project_name")
+    project_name = secure_filename(request.args.get("project_name"))
     pdf_dir = Config().get_pdfs_dir()
     pdf_path = os.path.join(pdf_dir, f"{project_name}.pdf")
 

src/project.py

@@ -144,3 +144,32 @@ def project_to_zip(self, project: str):
 
     def get_zip_path(self, project: str):
         return f"{self.get_project_path(project)}.zip"
+    
+    def get_project_files(self, project_name: str):
+        if not project_name:
+            return []
+
+        project_directory = "-".join(project_name.split(" "))
+        base_path = os.path.abspath(os.path.join(os.getcwd(), 'data', 'projects'))
+        directory = os.path.join(base_path, project_directory)
+
+        # Ensure the directory is within the allowed base path
+        if not os.path.exists(directory) or not os.path.commonprefix([directory, base_path]) == base_path:
+            return []
+
+        files = []
+        for root, _, filenames in os.walk(directory):
+            for filename in filenames:
+                file_relative_path = os.path.relpath(root, directory)
+                if file_relative_path == '.':
+                    file_relative_path = ''
+                file_path = os.path.join(file_relative_path, filename)
+                try:
+                    with open(os.path.join(root, filename), 'r') as file:
+                        files.append({
+                            "file": file_path,
+                            "code": file.read()
+                        })
+                except Exception as e:
+                    print(f"Error reading file {filename}: {e}")
+        return files

src/state.py

@@ -174,28 +174,4 @@ def get_latest_token_usage(self, project: str):
             if agent_state:
                 return json.loads(agent_state.state_stack_json)[-1]["token_usage"]
             return 0
-
-    def get_project_files(self, project_name: str):
-        if not project_name:
-            return []
-        project_directory = "-".join(project_name.split(" "))
-        directory = os.path.join(os.getcwd(), 'data', 'projects', project_directory) 
-        if(not os.path.exists(directory)):
-            return []
-        files = []
-        for root, _, filenames in os.walk(directory):
-            for filename in filenames:
-                file_relative_path = os.path.relpath(root, directory)
-                if file_relative_path == '.': file_relative_path = ''
-                file_path = os.path.join(file_relative_path, filename)
-                print("file_path",file_path)
-                try:
-                    with open(os.path.join(root, filename), 'r') as file:
-                        print("File:", filename)
-                        files.append({
-                            "file": file_path,
-                            "code": file.read()
-                        })
-                except Exception as e:
-                    print(f"Error reading file {filename}: {e}")
-        return files
+        

ui/bun.lockb

@@ -28,6 +28,7 @@
     "@xterm/xterm": "^5.5.0",
     "bits-ui": "^0.21.2",
     "clsx": "^2.1.0",
+    "dompurify": "^3.1.5",
     "mode-watcher": "^0.3.0",
     "paneforge": "^0.0.3",
     "socket.io-client": "^4.7.5",

ui/package.json

@@ -1,4 +1,5 @@
 <script>
+  import DOMPurify from "dompurify";
   import { emitMessage, socketListener } from "$lib/sockets";
   import { agentState, messages, isSending } from "$lib/store";
   import { calculateTokens } from "$lib/token";
@@ -11,12 +12,27 @@
     if (value !== null && value.agent_is_active == false) {
       isSending.set(false);
     }
-    if (value == null){
+    if (value == null) {
       inference_time = 0;
     }
   });
 
   let messageInput = "";
+
+  // Function to escape HTML
+  function escapeHTML(input) {
+    const map = {
+      "&": "&amp;",
+      "<": "&lt;",
+      ">": "&gt;",
+      '"': "&quot;",
+      "'": "&#039;",
+    };
+    return input.replace(/[&<>"']/g, function (m) {
+      return map[m];
+    });
+  }
+
   async function handleSendMessage() {
     const projectName = localStorage.getItem("selectedProject");
     const selectedModel = localStorage.getItem("selectedModel");
@@ -31,26 +47,29 @@
       return;
     }
 
-    if (messageInput.trim() !== "" && isSending) {
+    const sanitizedMessage = DOMPurify.sanitize(messageInput);
+    const escapedMessage = escapeHTML(sanitizedMessage);
+
+
+    if (messageInput.trim() !== "" && escapedMessage.trim() !== "" && isSending) {
       $isSending = true;
-      emitMessage("user-message", { 
-        message: messageInput,
+      emitMessage("user-message", {
+        message: escapedMessage,
         base_model: selectedModel,
         project_name: projectName,
         search_engine: serachEngine,
       });
       messageInput = "";
-      
     }
   }
   onMount(() => {
     socketListener("inference", function (data) {
-      if(data['type'] == 'time') {
+      if (data["type"] == "time") {
         inference_time = data["elapsed_time"];
       }
     });
   });
-       
+
   function setTokenSize(event) {
     const prompt = event.target.value;
     let tokens = calculateTokens(prompt);
@@ -73,41 +92,41 @@
       {/if}
     </div>
     <!-- {#if $agentState !== null} -->
-      <div class="px-1 rounded-md text-xs">
-        Model Inference: <span class="text-orange-600">{inference_time} sec</span>
-      </div>
+    <div class="px-1 rounded-md text-xs">
+      Model Inference: <span class="text-orange-600">{inference_time} sec</span>
+    </div>
     <!-- {/if} -->
   </div>
 
-<div class="expandable-input relative">
-  <textarea
-    id="message-input"
-    class="w-full p-4 font-medium focus:text-foreground rounded-xl outline-none h-28 pr-20 bg-secondary
-    {$isSending ? 'cursor-not-allowed' : ''}"   
-    placeholder="Type your message..."
-    disabled={$isSending}
-    bind:value={messageInput}
-    on:input={setTokenSize}
-    on:keydown={(e) => {
-      if (e.key === "Enter" && !e.shiftKey) {
-        e.preventDefault();
-        handleSendMessage();
-        document.querySelector('.token-count').textContent = 0;
-      }
-    }}
-  ></textarea>
-  <button 
-    on:click={handleSendMessage}
-    disabled={$isSending}
-    class="absolute text-secondary bg-primary p-2 right-4 bottom-6 rounded-full
+  <div class="expandable-input relative">
+    <textarea
+      id="message-input"
+      class="w-full p-4 font-medium focus:text-foreground rounded-xl outline-none h-28 pr-20 bg-secondary
     {$isSending ? 'cursor-not-allowed' : ''}"
-  >
-  {@html Icons.CornerDownLeft} 
-  </button>
-  <p class="absolute text-tertiary p-2 right-4 top-2">
-    <span class="token-count">0</span>
-  </p>
-</div>
+      placeholder="Type your message..."
+      disabled={$isSending}
+      bind:value={messageInput}
+      on:input={setTokenSize}
+      on:keydown={(e) => {
+        if (e.key === "Enter" && !e.shiftKey) {
+          e.preventDefault();
+          handleSendMessage();
+          document.querySelector(".token-count").textContent = 0;
+        }
+      }}
+    ></textarea>
+    <button
+      on:click={handleSendMessage}
+      disabled={$isSending}
+      class="absolute text-secondary bg-primary p-2 right-4 bottom-6 rounded-full
+    {$isSending ? 'cursor-not-allowed' : ''}"
+    >
+      {@html Icons.CornerDownLeft}
+    </button>
+    <p class="absolute text-tertiary p-2 right-4 top-2">
+      <span class="token-count">0</span>
+    </p>
+  </div>
 </div>
 
 <style>