only create parent dir when in memory transfer is false

CrowdStrike · Apr 1, 2024 · cb96318 · cb96318
1 parent 06876f4
commit cb96318
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 25 deletions.
diff --git a/falcon_data_replicator.py b/falcon_data_replicator.py
@@ -155,18 +155,6 @@ def handle_file(path, key, target_bkt, file_object=None, log_util: logging.Logge
 
 def download_message_files(msg, s3ta, s3or, log: logging.Logger):
     """Download the file specified in the SQS message and trigger file handling."""
-    # Construct output path for this message's files
-    msg_output_path = os.path.realpath(os.path.join(FDR.output_path, msg["pathPrefix"]))
-    # Only write files to the specified output_path
-    if os.path.commonpath([FDR.output_path, msg_output_path]) != FDR.output_path:
-        log.debug(
-            f"Skipping {msg_output_path} to prevent writes outside of output path: {FDR.output_path}"
-        )
-        return
-    # Ensure directory exists at output path
-    if not os.path.exists(msg_output_path):
-        # Create it if it doesn't
-        os.makedirs(msg_output_path)
     total_event_count = 0
     total_download_time_sec = 0.0
     total_transform_time_sec = 0.0
@@ -177,11 +165,23 @@ def download_message_files(msg, s3ta, s3or, log: logging.Logger):
         s3_path = s3_file['path']
         total_download_time_per_input_file = 0
         if not FDR.in_memory_transfer_only:
+            # Construct output path for this message's files
+            msg_output_path = os.path.realpath(os.path.join(FDR.output_path, msg["pathPrefix"]))
+            # Only write files to the specified output_path
+            if os.path.commonpath([FDR.output_path, msg_output_path]) != FDR.output_path:
+                log.info(
+                    f"Skipping {msg_output_path} to prevent writes outside of output path: {FDR.output_path}"
+                )
+                continue
+            # Ensure directory exists at output path
+            if not os.path.exists(msg_output_path):
+                # Create it if it doesn't
+                os.makedirs(msg_output_path)
             # Create a local path name for our destination file based off of the S3 path
             local_path = os.path.realpath(os.path.join(FDR.output_path, s3_path))
             # Only write files to the specified output_path
             if os.path.commonpath([FDR.output_path, local_path]) != FDR.output_path:
-                log.debug(
+                log.info(
                     f"Skipping {local_path} to prevent writes outside of output path: {FDR.output_path}"
                 )
                 continue

diff --git a/standalone/falcon_data_replicator.py b/standalone/falcon_data_replicator.py
@@ -156,24 +156,24 @@ def download_message_files(msg):
     download_message_files function will iterate through every file listed at msg['filePaths'],
     move it to our output_path, and then call handle_file.
     """
-    # Construct output path for this message's files
-    msg_output_path = os.path.realpath(os.path.join(FDR.output_path, msg["pathPrefix"]))
-    # Only write files to the specified output_path
-    if os.path.commonpath([FDR.output_path, msg_output_path]) != FDR.output_path:
-        logger.info(
-            f"Skipping {msg_output_path} to prevent writes outside of output path: {FDR.output_path}"
-        )
-        return
-    # Ensure directory exists at output path
-    if not os.path.exists(msg_output_path):
-        # Create it if it doesn't
-        os.makedirs(msg_output_path)
     # For every file in our message
     for s3_file in msg['files']:
         # Retrieve the bucket path for this file
         s3_path = s3_file['path']
         if not FDR.in_memory_transfer_only:
             # Create a local path name for our destination file based off of the S3 path
+            # Construct output path for this message's files
+            msg_output_path = os.path.realpath(os.path.join(FDR.output_path, msg["pathPrefix"]))
+            # Only write files to the specified output_path
+            if os.path.commonpath([FDR.output_path, msg_output_path]) != FDR.output_path:
+                logger.info(
+                    f"Skipping {msg_output_path} to prevent writes outside of output path: {FDR.output_path}"
+                )
+                continue 
+            # Ensure directory exists at output path
+            if not os.path.exists(msg_output_path):
+                # Create it if it doesn't
+                os.makedirs(msg_output_path)
             local_path = os.path.realpath(os.path.join(FDR.output_path, s3_path))
             # Only write files to the specified output_path
             if os.path.commonpath([FDR.output_path, local_path]) != FDR.output_path: