diff --git a/assemblyline_ui/api/v4/ui.py b/assemblyline_ui/api/v4/ui.py index 1bb69bdf..5cf9bce6 100644 --- a/assemblyline_ui/api/v4/ui.py +++ b/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 + + + --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) diff --git a/assemblyline_ui/config.py b/assemblyline_ui/config.py index 5d207c18..eae927b5 100644 --- a/assemblyline_ui/config.py +++ b/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) diff --git a/test/test_ui.py b/test/test_ui.py index b36ced92..e0f7359d 100644 --- a/test/test_ui.py +++ b/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!"