diff --git a/monkey/infection_monkey/exploit/i_agent_repository.py b/monkey/infection_monkey/exploit/i_agent_repository.py
index d825772a0..a00aac55a 100644
--- a/monkey/infection_monkey/exploit/i_agent_repository.py
+++ b/monkey/infection_monkey/exploit/i_agent_repository.py
@@ -1,6 +1,8 @@
 import abc
 import io
 
+from common import OperatingSystems
+
 # TODO: The Island also has an IAgentRepository with a totally different interface. At the moment,
 #       the Island and Agent have different needs, but at some point we should unify these.
 
@@ -13,10 +15,12 @@ class IAgentRepository(metaclass=abc.ABCMeta):
     """
 
     @abc.abstractmethod
-    def get_agent_binary(self, os: str, architecture: str = None) -> io.BytesIO:
+    def get_agent_binary(
+        self, operating_system: OperatingSystems, architecture: str = None
+    ) -> io.BytesIO:
         """
         Retrieve the appropriate agent binary from the repository.
-        :param str os: The name of the operating system on which the agent binary will run
+        :param operating_system: The name of the operating system on which the agent binary will run
         :param str architecture: Reserved
         :return: A file-like object for the requested agent binary
         :rtype: io.BytesIO