Merge pull request #10 from TrellixVulnTeam/master

CVE-2007-4559 Patch
jeromerony · Nov 25, 2022 · 787e863 · 787e863
2 parents f0dbc12 + 3414388
commit 787e863
Showing 1 changed file with 21 additions and 2 deletions.
diff --git a/prepare_data.py b/prepare_data.py
@@ -7,6 +7,25 @@
 from scipy.io import loadmat
 from torchvision.datasets.utils import download_url
 
+
+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) 
+
+
 ex1 = Experiment('Prepare CUB')
 
 
@@ -24,7 +43,7 @@ def download_extract_cub(cub_dir, cub_url):
     download_url(cub_url, root=path.dirname(cub_dir))
     filename = path.join(path.dirname(cub_dir), path.basename(cub_url))
     with tarfile.open(filename, 'r:gz') as tar:
-        tar.extractall(path=path.dirname(cub_dir))
+        safe_extract(tar, path=path.dirname(cub_dir))
 
 
 @ex1.capture
@@ -77,7 +96,7 @@ def download_extract_cars(cars_dir, cars_url, cars_annotations_url):
     download_url(cars_url, root=cars_dir)
     filename = path.join(cars_dir, path.basename(cars_url))
     with tarfile.open(filename, 'r:gz') as tar:
-        tar.extractall(path=cars_dir)
+        safe_extract(tar, path=cars_dir)
     return path.join(cars_dir, path.basename(cars_annotations_url))