diff --git a/monkey/infection_monkey/utils/propagation.py b/monkey/infection_monkey/utils/propagation.py
index 004bafdd2..9cfccab51 100644
--- a/monkey/infection_monkey/utils/propagation.py
+++ b/monkey/infection_monkey/utils/propagation.py
@@ -1,2 +1,5 @@
-def should_propagate(config: dict, current_depth: int) -> bool:
-    return config["config"]["depth"] > current_depth
+from common.configuration import AgentConfiguration
+
+
+def should_propagate(config: AgentConfiguration, current_depth: int) -> bool:
+    return config.propagation.maximum_depth > current_depth