From 22701fc0a3f17a5249b29c77ff0b586f15a97647 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Tue, 30 Aug 2022 06:46:35 -0400
Subject: [PATCH] Island: Redefine IAgentRepository

---
 .../monkey_island/cc/repository/__init__.py   |  2 +
 .../cc/repository/i_agent_repository.py       | 46 +++++++++++++++----
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py
index caef77b9b..e1d5fd47b 100644
--- a/monkey/monkey_island/cc/repository/__init__.py
+++ b/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
diff --git a/monkey/monkey_island/cc/repository/i_agent_repository.py b/monkey/monkey_island/cc/repository/i_agent_repository.py
index 5a784e5d2..c3b25dc34 100644
--- a/monkey/monkey_island/cc/repository/i_agent_repository.py
+++ b/monkey/monkey_island/cc/repository/i_agent_repository.py
@@ -1,15 +1,41 @@
-from abc import ABC
-from typing import Optional, Sequence
+from abc import ABC, abstractmethod
+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
-    def save_agent(self, agent: Monkey):
-        pass
+    """A repository used to store and retrieve `Agent` objects"""
 
-    def get_agents(
-        self, id: Optional[str] = None, running: Optional[bool] = None
-    ) -> Sequence[Monkey]:
-        pass
+    @abstractmethod
+    def upsert_agent(self, agent: Agent):
+        """
+        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`
+        """