Refactor API route for file upload from the UI to cleanup code + fix cases where NoneType is passed to os.path.exists

assemblyline_ui/api/v4/ui.py

@@ -2,14 +2,13 @@
 # UI ONLY APIs
 
 import os
+import tempfile
 
 from assemblyline.common import forge
 from assemblyline.common.bundling import import_bundle
-from assemblyline.common.str_utils import safe_str
-from assemblyline.common.uid import get_random_id
 from assemblyline.odm.messages.submission import Submission
 from assemblyline.odm.models.user import ROLES
-from assemblyline_core.submission_client import SubmissionClient, SubmissionException
+from assemblyline_core.submission_client import SubmissionClient
 from cart import get_metadata_only, is_cart
 from flask import request
 
@@ -63,10 +62,7 @@ def flowjs_check_chunk(**kwargs):
 
     Arguments (REQUIRED):
     flowChunkNumber      => Current chunk number
-    flowFilename         => Original filename
-    flowTotalChunks      => Total number of chunks
     flowIdentifier       => File unique identifier
-    flowCurrentChunkSize => Size of the current chunk
 
     Data Block:
     None
@@ -75,19 +71,11 @@ def flowjs_check_chunk(**kwargs):
     {'exists': True}     #Does the chunk exists on the server?
     """
 
-    flow_chunk_number = request.args.get("flowChunkNumber", None)
-    flow_chunk_size = request.args.get("flowChunkSize", None)
-    flow_total_size = request.args.get("flowTotalSize", None)
-    flow_filename = request.args.get("flowFilename", None)
-    flow_total_chunks = request.args.get("flowTotalChunks", None)
-    flow_identifier = request.args.get("flowIdentifier", None)
-    flow_current_chunk_size = request.args.get("flowCurrentChunkSize", None)
-
-    if not flow_chunk_number or not flow_identifier or not flow_current_chunk_size or not flow_filename \
-            or not flow_total_chunks or not flow_chunk_size or not flow_total_size:
-        return make_api_response("", "Required arguments missing. flowChunkNumber, flowIdentifier, "
-                                     "flowCurrentChunkSize, flowChunkSize and flowTotalSize "
-                                     "should always be present.", 412)
+    try:
+        flow_chunk_number = request.args["flowChunkNumber"]
+        flow_identifier = request.args["flowIdentifier"]
+    except KeyError as e:
+        return make_api_response("", f"Required argument missing: {e}", 412)
 
     filename = get_cache_name(flow_identifier, flow_chunk_number)
     with forge.get_cachestore("flowjs", config) as cache:
@@ -111,18 +99,26 @@ def flowjs_upload_chunk(**kwargs):
     Variables:
     None
 
-    Arguments (REQUIRED):
-    flowChunkNumber      => Current chunk number
-    flowChunkSize        => Usual size of the chunks
-    flowCurrentChunkSize => Size of the current chunk
-    flowTotalSize        => Total size for the file
-    flowIdentifier       => File unique identifier
-    flowFilename         => Original filename
-    flowRelativePath     => Relative path of the file on the client
-    flowTotalChunks      => Total number of chunks
+    Data Block (REQUIRED):
 
-    Data Block:
-    None
+    --0b34a3c50d3c02dd804a172329a0b2aa                          <-- Randomly generated boundary for this http request
+    Content-Disposition: form-data; name="flowChunkNumber"      <-- Current chunk number
+
+    1
+    --0b34a3c50d3c02dd804a172329a0b2aa                          <-- Switch to next part, file part
+    Content-Disposition: form-data; name="flowIdentifier"       <-- File unique identifier
+
+    00000000-0000-0000-0000-000000000000_000000000_testfile.txt
+    --0b34a3c50d3c02dd804a172329a0b2aa                          <-- Switch to next part, file part
+    Content-Disposition: form-data; name="flowTotalChunks"      <-- Total number of chunks
+
+    10
+    --0b34a3c50d3c02dd804a172329a0b2aa
+    Content-Disposition: form-data; name="file"; filename="testfile.txt" <-- File part
+    Content-Type: application/octet-stream
+
+    <BINARY DATA OF THE FILE TO UPLOAD... DOES NOT NEED TO BE ENCODED>
+    --0b34a3c50d3c02dd804a172329a0b2aa--                        <-- End of HTTP transmission
 
     Result example:
     {
@@ -131,27 +127,27 @@ def flowjs_upload_chunk(**kwargs):
      }
     """
 
-    flow_chunk_number = request.form.get("flowChunkNumber", None)
-    flow_chunk_size = request.form.get("flowChunkSize", None)
-    flow_current_chunk_size = request.form.get("flowCurrentChunkSize", None)
-    flow_total_size = request.form.get("flowTotalSize", None)
-    flow_identifier = request.form.get("flowIdentifier", None)
-    flow_filename = safe_str(request.form.get("flowFilename", None))
-    flow_relative_path = request.form.get("flowRelativePath", None)
-    flow_total_chunks = request.form.get("flowTotalChunks", None)
-    completed = True
-
-    if not flow_chunk_number or not flow_chunk_size or not flow_current_chunk_size or not flow_total_size \
-            or not flow_identifier or not flow_filename or not flow_relative_path or not flow_total_chunks:
-        return make_api_response("", "Required arguments missing. flowChunkNumber, flowChunkSize, "
-                                     "flowCurrentChunkSize, flowTotalSize, flowIdentifier, flowFilename, "
-                                     "flowRelativePath and flowTotalChunks should always be present.", 412)
-
-    filename = get_cache_name(flow_identifier, flow_chunk_number)
+    try:
+        flow_chunk_number = int(request.form["flowChunkNumber"])
+        flow_identifier = request.form["flowIdentifier"]
+        flow_total_chunks = int(request.form["flowTotalChunks"])
+    except KeyError as e:
+        return make_api_response("", f"Required argument missing: {e}", 412)
+    except ValueError as e:
+        return make_api_response("", f"Invalid argument type: {e}", 412)
+
+    # Evaluate if the chunk number is valid, if not return an error
+    if flow_chunk_number == 0 or flow_chunk_number > flow_total_chunks:
+        return make_api_response("", "Invalid chunk number", 412)
 
     with forge.get_cachestore("flowjs", config) as cache:
+        # Write the chunk to the cache
         file_obj = request.files['file']
-        cache.save(filename, file_obj.stream.read())
+        chunk_name = get_cache_name(flow_identifier, flow_chunk_number)
+        cache.save(chunk_name, file_obj.stream.read())
+
+        # Check if all chunks have been received, assume complete until proven otherwise
+        completed = True
 
         # Test in reverse order to fail fast
         for chunk in range(int(flow_total_chunks), 0, -1):
@@ -161,30 +157,21 @@ def flowjs_upload_chunk(**kwargs):
                 break
 
         if completed:
-            # Reconstruct the file
+            # Attempt file reconstruction and save to cache with the original identifier
             ui_sid = get_cache_name(flow_identifier)
-            target_file = os.path.join(TEMP_DIR, ui_sid)
-            try:
-                os.makedirs(TEMP_DIR)
-            except Exception:
-                pass
-
-            try:
-                os.unlink(target_file)
-            except Exception:
-                pass
-
-            for chunk in range(int(flow_total_chunks)):
-                chunk_name = get_cache_name(flow_identifier, chunk+1)
-                with open(target_file, "ab") as t:
-                    t.write(cache.get(chunk_name))
-                cache.delete(chunk_name)
-
-            # Save the reconstructed file
-            with open(target_file, "rb") as t:
-                cache.save(ui_sid, t.read())
-
-            os.unlink(target_file)
+            with tempfile.NamedTemporaryFile(dir=TEMP_DIR) as target_file:
+                # Iterate through the chunks in order to reconstruct the file
+                for chunk in range(int(flow_total_chunks)):
+                    chunk_name = get_cache_name(flow_identifier, chunk+1)
+                    # Write chunks to temporary file
+                    target_file.write(cache.get(chunk_name))
+                    # Delete the chunk from the cache
+                    cache.delete(chunk_name)
+
+                # Once the file is reconstructed, save it to the cache with the original identifier
+                target_file.flush()
+                target_file.seek(0)
+                cache.save(ui_sid, target_file.read())
 
     return make_api_response({'success': True, 'completed': completed})
 
@@ -223,82 +210,66 @@ def start_ui_submission(ui_sid, **kwargs):
     ui_params = request.json
     submit_result = None
     submitted_file = None
-    target_dir = None
 
-    try:
-        # Download the file from the cache
-        with forge.get_cachestore("flowjs", config) as cache:
-            ui_sid = get_cache_name(ui_sid)
-            fname = ui_params.pop('filename', ui_sid)
-            if cache.exists(ui_sid):
-                target_dir = os.path.join(TEMP_DIR, ui_sid)
-                os.makedirs(target_dir, exist_ok=True)
-
-                target_file = os.path.join(target_dir, get_random_id())
-
-                if os.path.exists(target_file):
-                    os.unlink(target_file)
-
-                # Save the reconstructed file
-                cache.download(ui_sid, target_file)
-                submitted_file = target_file
-
-        # Submit the file
-        if submitted_file is not None:
-            with open(submitted_file, 'rb') as fh:
-                if is_cart(fh.read(256)):
+    # Download the file from the cache
+
+    with forge.get_cachestore("flowjs", config) as cache:
+        ui_sid = get_cache_name(ui_sid)
+        fname = ui_params.pop('filename', ui_sid)
+        if not cache.exists(ui_sid):
+            # No file was found for the given ID, return an error and decrement the submission quota for the user
+            decrement_submission_quota(user)
+            return make_api_response({"started": False, "sid": None}, "No files where found for ID %s. "
+                                                                        "Try again..." % ui_sid, 404)
+        # Submit to dispatcher
+        try:
+            # Save uploaded file to a temporary file for processing
+            submitted_file = tempfile.NamedTemporaryFile(dir=TEMP_DIR).name
+            cache.download(ui_sid, submitted_file)
+
+            # Check if the file is a CART bundle
+            with open(submitted_file, "rb") as temp_file:
+                if is_cart(temp_file.read(256)):
                     meta = get_metadata_only(submitted_file)
                     if meta.get('al', {}).get('type', 'unknown') == 'archive/bundle/al':
-                        try:
-                            submission = import_bundle(submitted_file, allow_incomplete=True, identify=IDENTIFY,
-                                                       dtl=ui_params.get('ui_params', {}).get('ttl'))
-                        except Exception as e:
-                            return make_api_response("", err=str(e), status_code=400)
+                        # Import the submission bundle and return the submission ID
+                        submission = import_bundle(submitted_file, allow_incomplete=True, identify=IDENTIFY,
+                                                   dtl=ui_params.get('ui_params', {}).get('ttl'))
                         return make_api_response({"started": True, "sid": submission['sid']})
 
-            # Submit to dispatcher
-            try:
-                # Initialize submission validation process
-                _, _, _, _, s_params, metadata = init_submission(request, user, endpoint="ui")
-                s_params['quota_item'] = True
-                allow_description_overwrite = False
-                if not s_params.get("description"):
-                    # If no custom description is specified, create one based on filename
-                    s_params["description"] = f"Inspection of file: {fname}"
-                    allow_description_overwrite = True
-                submission_obj = Submission({
-                    "files": [],
-                    "metadata": metadata,
-                    "params": s_params
-                })
-            except (ValueError, KeyError) as e:
-                return make_api_response("", err=str(e), status_code=400)
-
-            try:
-                submit_result = SubmissionClient(
-                    datastore=STORAGE, filestore=FILESTORE, config=config, identify=IDENTIFY
-                ).submit(
-                    submission_obj,
-                    local_files=[(fname, submitted_file)],
-                    allow_description_overwrite=allow_description_overwrite
-                )
-                submission_received(submission_obj)
-            except SubmissionException as e:
-                return make_api_response("", err=str(e), status_code=400)
-
+            # Initialize submission validation process
+            _, _, _, _, s_params, metadata = init_submission(request, user, endpoint="ui")
+            s_params['quota_item'] = True
+            allow_description_overwrite = False
+            if not s_params.get("description"):
+                # If no custom description is specified, create one based on filename
+                s_params["description"] = f"Inspection of file: {fname}"
+                allow_description_overwrite = True
+            submission_obj = Submission({
+                "files": [],
+                "metadata": metadata,
+                "params": s_params
+            })
+
+            # Attempt to submit the file to the dispatcher
+            submit_result = SubmissionClient(
+                datastore=STORAGE, filestore=FILESTORE, config=config, identify=IDENTIFY
+            ).submit(
+                submission_obj,
+                local_files=[(fname, submitted_file)],
+                allow_description_overwrite=allow_description_overwrite
+            )
+            submission_received(submission_obj)
             return make_api_response({"started": True, "sid": submit_result.sid})
-        else:
-            return make_api_response({"started": False, "sid": None}, "No files where found for ID %s. "
-                                                                      "Try again..." % ui_sid, 404)
-    finally:
-        if submit_result is None:
-            # We had an error during the submission, release the quotas for the user
-            decrement_submission_quota(user)
-
-        # Remove file
-        if os.path.exists(submitted_file):
-            os.unlink(submitted_file)
-
-        # Remove dir
-        if os.path.exists(target_dir) and os.path.isdir(target_dir):
-            os.rmdir(target_dir)
+        except Exception as e:
+            # If an error occurs during submission, return an error response
+            return make_api_response("", err=str(e), status_code=400)
+        finally:
+            # Perform cleanup actions regardless of whether submission was successful or not
+            if submit_result is None:
+                # We had an error during the submission, release the quotas for the user
+                decrement_submission_quota(user)
+
+            if submitted_file and os.path.exists(submitted_file):
+                # Clean up the temporary file
+                os.remove(submitted_file)

assemblyline_ui/config.py

@@ -73,6 +73,11 @@ def load_openid_configuration(old_config: OAuthProvider, url: str) -> OAuthProvi
 TEMP_DIR = "/var/lib/assemblyline/flowjs/"
 TEMP_SUBMIT_DIR = "/var/lib/assemblyline/submit/"
 
+# Initialize all temporary directories used by the UI (mostly for testing purposes since Dockerfile creates them)
+for temp_dir in [TEMP_DIR, TEMP_SUBMIT_DIR, BUNDLING_DIR]:
+    if not os.path.exists(temp_dir):
+        os.makedirs(temp_dir, exist_ok=True)
+
 redis_persistent = get_client(config.core.redis.persistent.host, config.core.redis.persistent.port, False)
 redis = get_client(config.core.redis.nonpersistent.host, config.core.redis.nonpersistent.port, False)
 

test/test_ui.py

@@ -47,9 +47,10 @@ def upload_file_flowjs(session, host, data=None):
     chunk_size = math.ceil(len(data) / total_chunk)
     ui_id = get_random_id()
 
-    counter = 0
-    while True:
-        x = counter % total_chunk
+    # Purposefully shuffle the chunk order to ensure that the server is able to handle out of order chunks
+    chunk_parts = list(range(total_chunk))
+    random.shuffle(chunk_parts)
+    for x in chunk_parts:
         try:
             params = {
                 "flowChunkNumber": f"{x + 1}",
@@ -61,8 +62,6 @@ def upload_file_flowjs(session, host, data=None):
                 "flowCurrentChunkSize": f"{chunk_size}",
             }
             resp = get_api_data(session, f"{host}/api/v4/ui/flowjs/", params=params)
-            if resp["exist"]:
-                counter += 1
         except APIError as e:
             if (
                 str(e) == "Chunk does not exist, please send it!"