Island: Redefine IAgentRepository

monkey/monkey_island/cc/repository/__init__.py

@@ -8,6 +8,7 @@ from .i_simulation_repository import ISimulationRepository
 from .i_credentials_repository import ICredentialsRepository
 from .i_user_repository import IUserRepository
 from .i_machine_repository import IMachineRepository
+from .i_agent_repository import IAgentRepository
 
 
 from .local_storage_file_repository import LocalStorageFileRepository
@@ -21,3 +22,4 @@ from .file_simulation_repository import FileSimulationRepository
 from .json_file_user_repository import JSONFileUserRepository
 from .mongo_credentials_repository import MongoCredentialsRepository
 from .mongo_machine_repository import MongoMachineRepository
+from .mongo_agent_repository import MongoAgentRepository

monkey/monkey_island/cc/repository/i_agent_repository.py

@@ -1,15 +1,41 @@
-from abc import ABC
+from abc import ABC, abstractmethod
-from typing import Optional, Sequence
+from typing import Sequence
 
-from monkey_island.cc.models import Monkey
+from monkey_island.cc.models import Agent, AgentID
 
 
 class IAgentRepository(ABC):
-    # TODO rename Monkey document to Agent
+    """A repository used to store and retrieve `Agent` objects"""
-    def save_agent(self, agent: Monkey):
-        pass
 
-    def get_agents(
+    @abstractmethod
-        self, id: Optional[str] = None, running: Optional[bool] = None
+    def upsert_agent(self, agent: Agent):
-    ) -> Sequence[Monkey]:
+        """
-        pass
+        Upsert (insert or update) an `Agent`
+
+        Insert the `Agent` if no `Agent` with a matching ID exists in the repository. If the agent
+        already exists, update it.
+
+        :param agent: The `agent` to be inserted or updated
+        :raises StorageError: If an error occurred while attempting to store the `Agent`
+        """
+
+    @abstractmethod
+    def get_agent_by_id(self, id: AgentID) -> Agent:
+        """
+        Get an `Agent` by ID
+
+        :param id: The ID of the `Agent` to be retrieved
+        :return: An `Agent` with a matching `id`
+        :raises UnknownRecordError: If an `Agent` with the specified `id` does not exist in the
+                                    repository
+        :raises RetrievalError: If an error occurred while attempting to retrieve the `Agent`
+        """
+
+    @abstractmethod
+    def get_running_agents(self) -> Sequence[Agent]:
+        """
+        Get all `Agents` that are currently running
+
+        :return: All `Agents` that are currently running
+        :raises RetrievalError: If an error occurred while attempting to retrieve the `Agents`
+        """