CVE-2007-4559 Patch

mobilenet/dataset/scripts/tar_to_record.py

@@ -22,7 +22,29 @@ def main(dataset_paths, percent_eval, directory):
         sys.exit(1)
     for dataset in dataset_paths:
         with tarfile.open(dataset) as tar_file:
-            tar_file.extractall(join(EXTRACT_PATH, 'out'))
+
+            import os
+
+            def is_within_directory(directory, target):
+
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                return prefix == abs_directory
+
+            def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+
+                tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+            safe_extract(tar_file, join(EXTRACT_PATH,"out"))
 
     if percent_eval > 100 or percent_eval < 0:
         percent_eval = 30

mobilenet/tflite/scripts/prepare_checkpoint.py

@@ -12,4 +12,26 @@ def prepare_checkpoint(checkpoint_tar):
     model_name = "/opt/ml/model/ssd_mobilenet_v2_quantized_300x300_coco_2019_01_03" + '.tar.gz'
 
     with tarfile.open(model_name) as model:
-        model.extractall(CKPT_DIR)
+
+        import os
+
+        def is_within_directory(directory, target):
+
+            abs_directory = os.path.abspath(directory)
+            abs_target = os.path.abspath(target)
+
+            prefix = os.path.commonprefix([abs_directory, abs_target])
+
+            return prefix == abs_directory
+
+        def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+            for member in tar.getmembers():
+                member_path = os.path.join(path, member.name)
+                if not is_within_directory(path, member_path):
+                    raise Exception("Attempted Path Traversal in Tar File")
+
+            tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+        safe_extract(model, CKPT_DIR)