@@ -59,6 +59,9 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
5959
6060
6161def 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+
6265 for module_path in possible_module_paths :
6366 try :
6467 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
7174 except Exception :
7275 logger .debug (f"Patchflow { patchflow } not found as a file/directory in { module_path } " )
7376
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 } " )
8286
8387 return None
8488
0 commit comments