improve: security - restriected_getattr

__init__.py

@@ -8,7 +8,7 @@
 import importlib
 import logging
 
-version_code = [1, 2, 1]
+version_code = [1, 2, 2]
 version_str = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 logging.info(f"### Loading: ComfyUI-Impact-Subpack ({version_str})")
 

modules/subcore.py

@@ -37,6 +37,18 @@ def create_segmasks(results):
     return results
 
 
+# Limit the commands that can be executed through `getattr` to `ultralytics.nn.modules.head.Detect.forward`.
+def restricted_getattr(obj, name, *args):
+    if name != "forward":
+        logging.error(f"Access to potentially dangerous attribute '{obj.__module__}.{obj.__name__}.{name}' is blocked.\nIf you believe the use of this code is genuinely safe, please report it.\nhttps://github.com/ltdrdata/ComfyUI-Impact-Subpack/issues")
+        raise RuntimeError(f"Access to potentially dangerous attribute '{obj.__module__}.{obj.__name__}.{name}' is blocked.")
+        
+    return getattr(obj, name, *args)
+
+restricted_getattr.__module__ = 'builtins'
+restricted_getattr.__name__ = 'getattr'
+
+
 try:
     from ultralytics import YOLO
     from ultralytics.nn.tasks import DetectionModel
@@ -108,7 +120,7 @@ try:
 
         torch_whitelist += [DetectionModel, aliasYOLOv10DetectionModel, SegmentationModel, IterableSimpleNamespace,
                             aliasIterableSimpleNamespace, TaskAlignedAssigner, aliasTaskAlignedAssigner, aliasv10DetectLoss,
-                            getattr, dill._dill._load_type, scalar, dtype, Float64DType]
+                            restricted_getattr, dill._dill._load_type, scalar, dtype, Float64DType]
 
     build_torch_whitelist()
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-impact-subpack"
 description = "This node pack provides nodes that complement the Impact Pack, such as the UltralyticsDetectorProvider."
-version = "1.2.1"
+version = "1.2.2"
 license = { file = "LICENSE.txt" }
 dependencies = ["matplotlib", "ultralytics==8.3.40"]