From ac97d4bbaa01d4f4e1327477539c188a46c05071 Mon Sep 17 00:00:00 2001
From: Daniel Markstedt <markstedt@gmail.com>
Date: Mon, 20 Sep 2021 12:47:39 -0700
Subject: [PATCH] Use save() to upload files, to address file corruption issues
 with the previous solution. Also adds file ending validation which has been
 requested by users.

---
 src/web/templates/index.html | 14 +++-------
 src/web/web.py               | 52 +++++++++++++++++++++++-------------
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/src/web/templates/index.html b/src/web/templates/index.html
index 6cc704d9..549253b2 100644
--- a/src/web/templates/index.html
+++ b/src/web/templates/index.html
@@ -170,22 +170,16 @@
     <h2>Upload File</h2>
     <p>Uploads file to <tt>{{base_dir}}</tt>. Max file size is set to {{max_file_size / 1024 /1024 }}MB</p>
     <table style="border: none">
-            <tr style="border: none">
-                <td style="border: none; vertical-align:top;">
-                <form id="uploadForm" action="/files/upload/" onchange="fileSelect(event)" method="post" enctype="multipart/form-data">
+        <tr style="border: none">
+            <td style="border: none; vertical-align:top;">
+		<form action="/files/upload" method="post" enctype="multipart/form-data">
                     <label for="file">File:</label>
-                    <input type="file" name="file"/>
+                    <input type="file" name="file" />
                     <input type="submit" value="Upload" />
                 </form>
             </td>
         </tr>
     </table>
-<script>
-    function fileSelect(e) {
-        document.getElementById("uploadForm").setAttribute('action', "/files/upload/" + e.target.files[0].name)
-        console.log(e.target.files[0].name);
-    }
-</script>
 
     <hr/>
 
diff --git a/src/web/web.py b/src/web/web.py
index c2dec75e..cdcb86a4 100644
--- a/src/web/web.py
+++ b/src/web/web.py
@@ -1,4 +1,13 @@
-from flask import Flask, render_template, request, flash, url_for, redirect, send_file, send_from_directory
+from flask import (
+    Flask,
+    render_template,
+    request,
+    flash,
+    url_for,
+    redirect,
+    send_file,
+    send_from_directory,
+)
 
 from file_cmds import (
     list_files,
@@ -465,29 +474,34 @@ def download_img():
         return redirect(url_for("index"))
 
 
-@app.route("/files/upload/<filename>", methods=["POST"])
-def upload_file(filename):
-    if not filename:
-        flash("No file provided.", "error")
+@app.route("/files/upload", methods=["POST"])
+def upload_file():
+    if 'file' not in request.files:
+        flash("No file part in request.", "error")
+        return redirect(url_for("index"))
+    f = request.files["file"]
+    if f.filename == "":
+        flash("No file selected.", "error")
         return redirect(url_for("index"))
 
+    from werkzeug.utils import secure_filename
     from os import path
-    file_path = path.join(app.config["UPLOAD_FOLDER"], filename)
-    if path.isfile(file_path):
+    filename = secure_filename(f.filename)
+    filepath = path.join(app.config["UPLOAD_FOLDER"], filename)
+    if not filename.lower().endswith(VALID_FILE_SUFFIX):
+        flash("Not a file format RaSCSI recognizes. Needs to be one of:", "error")
+        flash(f"{VALID_FILE_SUFFIX}", "error")
+        return redirect(url_for("index"))
+    if path.isfile(filepath):
         flash(f"{filename} already exists.", "error")
         return redirect(url_for("index"))
-
-    from io import DEFAULT_BUFFER_SIZE
-    binary_new_file = "bx"
-    with open(file_path, binary_new_file, buffering=DEFAULT_BUFFER_SIZE) as f:
-        chunk_size = DEFAULT_BUFFER_SIZE
-        while True:
-            chunk = request.stream.read(chunk_size)
-            if len(chunk) == 0:
-                break
-            f.write(chunk)
-    # TODO: display an informative success message
-    return redirect(url_for("index", filename=filename))
+    try:
+        f.save(filepath)
+        flash(f"File {filename} successfully uploaded to {base_dir} !")
+        return redirect(url_for("index"))
+    except error as e:
+        flash(f"Failed to upload {filename}: {str(e)}")
+        return redirect(url_for("index"))
 
 
 @app.route("/files/create", methods=["POST"])