@@ -59,6 +59,9 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
59
59
60
60
61
61
def find_patchflow (possible_module_paths : Iterable [str ], patchflow : str ) -> Any | None :
62
+ # Define a whitelist of allowed module paths
63
+ allowed_modules = {"allowed_module1" , "allowed_module2" , "allowed_module3" }
64
+
62
65
for module_path in possible_module_paths :
63
66
try :
64
67
spec = importlib .util .spec_from_file_location ("custom_module" , module_path )
@@ -71,14 +74,15 @@ def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any
71
74
except Exception :
72
75
logger .debug (f"Patchflow { patchflow } not found as a file/directory in { module_path } " )
73
76
74
- try :
75
- module = importlib .import_module (module_path )
76
- logger .info (f"Patchflow { patchflow } loaded from { module_path } " )
77
- return getattr (module , patchflow )
78
- except ModuleNotFoundError :
79
- logger .debug (f"Patchflow { patchflow } not found as a module in { module_path } " )
80
- except AttributeError :
81
- logger .debug (f"Patchflow { patchflow } not found in { module_path } " )
77
+ if module_path in allowed_modules :
78
+ try :
79
+ module = importlib .import_module (module_path )
80
+ logger .info (f"Patchflow { patchflow } loaded from { module_path } " )
81
+ return getattr (module , patchflow )
82
+ except ModuleNotFoundError :
83
+ logger .debug (f"Patchflow { patchflow } not found as a module in { module_path } " )
84
+ except AttributeError :
85
+ logger .debug (f"Patchflow { patchflow } not found in { module_path } " )
82
86
83
87
return None
84
88
0 commit comments