From 20b79298bebd49f4bfead4105ed954abadb05734 Mon Sep 17 00:00:00 2001
From: TrellixVulnTeam <charles.mcfarland@trellix.com>
Date: Tue, 18 Oct 2022 01:39:59 +0000
Subject: [PATCH] Adding tarfile member sanitization to extractall()

---
 scripts/coxph_evaluation.py | 21 ++++++++++++++++++++-
 scripts/evaluation.py       | 21 ++++++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/scripts/coxph_evaluation.py b/scripts/coxph_evaluation.py
index f0c6c90..1803720 100644
--- a/scripts/coxph_evaluation.py
+++ b/scripts/coxph_evaluation.py
@@ -64,7 +64,26 @@ def main(args):
     logger.info(f"Extracting model from path: {model_path}")
 
     with tarfile.open(model_path) as tar:
-        tar.extractall(path=".")
+        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, path=".")
     logger.info("Loading model")
     with open(args.model_name, "rb") as f:
         model = pickle.load(f)
diff --git a/scripts/evaluation.py b/scripts/evaluation.py
index 96b006c..b95802b 100644
--- a/scripts/evaluation.py
+++ b/scripts/evaluation.py
@@ -57,7 +57,26 @@ def main(args):
     logger.info(f"Extracting model from path: {model_path}")
 
     with tarfile.open(model_path) as tar:
-        tar.extractall(path=".")
+        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, path=".")
     logger.info("Loading model")
     with open(args.model_name, "rb") as f:
         model = pickle.load(f)