diff --git a/monkey/infection_monkey/master/option_parsing.py b/monkey/infection_monkey/master/option_parsing.py
index c35bf6303..c9262c5c9 100644
--- a/monkey/infection_monkey/master/option_parsing.py
+++ b/monkey/infection_monkey/master/option_parsing.py
@@ -1,13 +1,12 @@
-from typing import Dict
-
+from common.configuration import CustomPBAConfiguration
 from infection_monkey.utils.environment import is_windows_os
 
 
-def custom_pba_is_enabled(pba_options: Dict) -> bool:
+def custom_pba_is_enabled(pba_options: CustomPBAConfiguration) -> bool:
     if not is_windows_os():
-        if pba_options["linux_command"]:
+        if pba_options.linux_command:
             return True
     else:
-        if pba_options["windows_command"]:
+        if pba_options.windows_command:
             return True
     return False