Command injection and Path traversal security fixes (#888)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/data/dataloaders/v5loader.py

@@ -54,7 +54,7 @@ for orientation in ExifTags.TAGS.keys():
 def get_hash(paths):
     # Returns a single hash value of a list of paths (files or dirs)
     size = sum(os.path.getsize(p) for p in paths if os.path.exists(p))  # sizes
-    h = hashlib.md5(str(size).encode())  # hash sizes
+    h = hashlib.sha256(str(size).encode())  # hash sizes
     h.update(''.join(paths).encode())  # hash paths
     return h.hexdigest()  # return hash
 

ultralytics/yolo/data/utils.py

@@ -44,7 +44,7 @@ def img2label_paths(img_paths):
 def get_hash(paths):
     # Returns a single hash value of a list of paths (files or dirs)
     size = sum(os.path.getsize(p) for p in paths if os.path.exists(p))  # sizes
-    h = hashlib.md5(str(size).encode())  # hash sizes
+    h = hashlib.sha256(str(size).encode())  # hash sizes
     h.update("".join(paths).encode())  # hash paths
     return h.hexdigest()  # return hash
 

ultralytics/yolo/engine/trainer.py

@@ -5,6 +5,7 @@ Simple training loop; Boilerplate that could apply to any arbitrary neural netwo
 
 import os
 import subprocess
+import sys
 import time
 from collections import defaultdict
 from copy import deepcopy
@@ -28,10 +29,10 @@ from ultralytics.yolo.utils import (DEFAULT_CFG, LOGGER, RANK, SETTINGS, TQDM_BA
                                     yaml_save)
 from ultralytics.yolo.utils.autobatch import check_train_batch_size
 from ultralytics.yolo.utils.checks import check_file, check_imgsz, print_args
-from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_command
+from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_file, find_free_network_port
 from ultralytics.yolo.utils.files import get_latest_run, increment_path
 from ultralytics.yolo.utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, init_seeds, one_cycle,
-                                                select_device, strip_optimizer)
+                                                select_device, strip_optimizer, TORCH_1_9)
 
 
 class BaseTrainer:
@@ -174,13 +175,18 @@ class BaseTrainer:
 
         # Run subprocess if DDP training, else train normally
         if world_size > 1 and "LOCAL_RANK" not in os.environ:
-            command = generate_ddp_command(world_size, self)
+            # cmd, file = generate_ddp_command(world_size, self)  # security vulnerability in Snyk scans
+            file = generate_ddp_file(self) if sys.argv[0].endswith('yolo') else os.path.abspath(sys.argv[0])
+            torch_distributed_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch"
+            cmd = [
+                sys.executable, "-m", torch_distributed_cmd, "--nproc_per_node", f"{world_size}", "--master_port",
+                f"{find_free_network_port()}", file] + sys.argv[1:]
             try:
-                subprocess.run(command)
+                subprocess.run(cmd, check=True)
             except Exception as e:
-                self.console(e)
+                self.console.warning(e)
             finally:
-                ddp_cleanup(command, self)
+                ddp_cleanup(self, file)
         else:
             self._do_train(int(os.getenv("RANK", -1)), world_size)
 

ultralytics/yolo/utils/dist.py

@@ -44,21 +44,15 @@ def generate_ddp_file(trainer):
 
 def generate_ddp_command(world_size, trainer):
     import __main__  # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
-    file_name = os.path.abspath(sys.argv[0])
+    file = generate_ddp_file(trainer) if sys.argv[0].endswith('yolo') else os.path.abspath(sys.argv[0])
-    using_cli = not file_name.endswith(".py")
-    if using_cli:
-        file_name = generate_ddp_file(trainer)
     torch_distributed_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch"
-    return [
+    cmd = [
         sys.executable, "-m", torch_distributed_cmd, "--nproc_per_node", f"{world_size}", "--master_port",
-        f"{find_free_network_port()}", file_name] + sys.argv[1:]
+        f"{find_free_network_port()}", file] + sys.argv[1:]
+    return cmd, file
 
 
-def ddp_cleanup(command, trainer):
+def ddp_cleanup(trainer, file):
     # delete temp file if created
-    tempfile_suffix = f"{id(trainer)}.py"
+    if f"{id(trainer)}.py" in file:  # if temp_file suffix in file
-    if tempfile_suffix in "".join(command):
+        os.remove(file)
-        for chunk in command:
-            if tempfile_suffix in chunk:
-                os.remove(chunk)
-                break