diff --git a/donut/algorithm.py b/donut/algorithm.py index 94b88d1..3cbbdc2 100644 --- a/donut/algorithm.py +++ b/donut/algorithm.py @@ -83,7 +83,29 @@ def save_model(model_vs, args: AlgorithmArgs): def load_model(model_vs, args: AlgorithmArgs): # decompress archive with model files with tarfile.open(args.modelInput, "r:gz") as f: - f.extractall() + + 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(f) saver = VariableSaver(get_variables_as_dict(model_vs), MODEL_PATH) saver.restore() diff --git a/omnianomaly/algorithm.py b/omnianomaly/algorithm.py index 049f406..e0831fd 100644 --- a/omnianomaly/algorithm.py +++ b/omnianomaly/algorithm.py @@ -110,7 +110,29 @@ def save(self): def load(self): # decompress archive with model files with tarfile.open(self.args.modelInput, "r:gz") as f: - f.extractall() + + 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(f) saver = VariableSaver(get_variables_as_dict(self.model_vs), self.model_dir) saver.restore() diff --git a/torsk/algorithm.py b/torsk/algorithm.py index 2fb84dc..3e4c722 100644 --- a/torsk/algorithm.py +++ b/torsk/algorithm.py @@ -120,7 +120,29 @@ def save_model(model, outfile: Path) -> None: def load_model(config: AlgorithmArgs) -> Any: # decompress archive to folder with tarfile.open(config.modelInput, "r:gz") as f: - f.extractall() + + 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(f) return torsk.load_model(MODEL_PATH)