From eb3fe21b11d93f61a240946c9fc881ec1722dee1 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 29 Aug 2022 15:34:36 -0400
Subject: [PATCH] Island: Redefine IMachineRepository

---
 .../cc/repository/i_machine_repository.py     | 80 +++++++++++++++----
 vulture_allowlist.py                          |  7 +-
 2 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py
index 79538e22f..97bd3cf87 100644
--- a/monkey/monkey_island/cc/repository/i_machine_repository.py
+++ b/monkey/monkey_island/cc/repository/i_machine_repository.py
@@ -1,19 +1,71 @@
-from abc import ABC
-from typing import Optional, Sequence
+from abc import ABC, abstractmethod
+from ipaddress import IPv4Address
+from typing import Sequence
 
-from monkey_island.cc.models import Machine
+from common.types import HardwareID
+from monkey_island.cc.models import Machine, MachineID
 
 
 class IMachineRepository(ABC):
-    def save_machine(self, machine: Machine):
-        pass
+    """A repository used to store and retrieve Machines"""
 
-    # TODO define or re-use machine state.
-    def get_machines(
-        self,
-        id: Optional[str] = None,
-        ips: Optional[Sequence[str]] = None,
-        state: Optional[MachineState] = None,  # noqa: F821
-        is_island: Optional[bool] = None,
-    ) -> Sequence[Machine]:
-        pass
+    @abstractmethod
+    def create_machine(self) -> Machine:
+        """
+        Create a new `Machine` in the repository
+
+        :return: A new `Machine` with a unique ID
+        :raises StorageError: If a new `Machine` could not be created
+        """
+
+    @abstractmethod
+    def update_machine(self, machine: Machine):
+        """
+        Update an existing `Machine` in the repository
+
+        :param machine: An updated Machine object to store in the repository
+        :raises UnknownRecordError: If the provided `Machine` does not exist in the repository
+        :raises StorageError: If an error occurred while attempting to store the `Machine`
+        """
+
+    @abstractmethod
+    def get_machine_by_id(self, id_: MachineID) -> Machine:
+        """
+        Get a `Machine` by ID
+
+        :param id: The ID of the `Machine` to be retrieved
+        :return: A `Machine` with a matching `id`
+        :raises UnknownRecordError: If a `Machine` with the specified `id` does not exist in the
+                                    repository
+        :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine`
+        """
+
+    @abstractmethod
+    def get_machine_by_hardware_id(self, hardware_id: HardwareID) -> Machine:
+        """
+        Get a `Machine` by ID
+
+        :param hardware_id: The hardware ID of the `Machine` to be retrieved
+        :return: A `Machine` with a matching `hardware_id`
+        :raises UnknownRecordError: If a `Machine` with the specified `hardware_id` does not exist
+                                    in the repository
+        :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine`
+        """
+
+    @abstractmethod
+    def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]:
+        """
+        Search for machines by IP address
+
+        :param ip: The IP address to search for
+        :return: A sequence of Machines that have a network interface with a matching IP
+        :raises UnknownRecordError: If a `Machine` with the specified `ip` does not exist in the
+                                    repository
+        :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine`
+        """
+
+    @abstractmethod
+    def reset(self):
+        """
+        Removes all data from the repository
+        """
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
index b203a2bea..9da5fee84 100644
--- a/vulture_allowlist.py
+++ b/vulture_allowlist.py
@@ -244,8 +244,11 @@ IConfigRepository.get_config_field
 ILogRepository.get_logs
 ILogRepository.save_log
 ILogRepository.delete_log
-IMachineRepository.save_machine
-IMachineRepository.get_machines
+IMachineRepository.create_machine
+IMachineRepository.update_machine
+IMachineRepository.get_machine_by_id
+IMachineRepository.get_machine_by_hardware_id
+IMachineRepository.get_machines_by_ip
 INetworkMapRepository.get_map
 INetworkMapRepository.save_netmap
 IReportRepository