3078 otp repository

monkey/monkey_island/cc/repositories/__init__.py

@@ -11,6 +11,7 @@
 from .i_agent_event_repository import IAgentEventRepository
 from .i_agent_log_repository import IAgentLogRepository
 from .i_agent_plugin_repository import IAgentPluginRepository
+from .i_otp_repository import IOTPRepository
 
 
 from .local_storage_file_repository import LocalStorageFileRepository
@@ -29,6 +30,7 @@
 from .mongo_agent_repository import MongoAgentRepository
 from .mongo_node_repository import MongoNodeRepository
 from .mongo_agent_event_repository import MongoAgentEventRepository
+from .mongo_otp_repository import MongoOTPRepository
 from .file_agent_log_repository import FileAgentLogRepository
 from .file_agent_plugin_repository import FileAgentPluginRepository
 

monkey/monkey_island/cc/repositories/i_otp_repository.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+
+
+class IOTPRepository(ABC):
+    """A repository used to store and retrieve `OTP`s"""
+
+    @abstractmethod
+    def insert_otp(self, otp: str, expiration: float):
+        """
+        Insert an OTP into the repository
+
+        :param otp: The OTP to insert
+        :param expiration: The time that the OTP expires
+        :raises StorageError: If an error occurs while attempting to insert the OTP
+        """
+
+    @abstractmethod
+    def get_expiration(self, otp: str) -> float:
+        """
+        Get the expiration time of a given OTP
+
+        :param otp: OTP for which to get the expiration time
+        :return: The time that the OTP expires
+        :raises RetrievalError: If an error occurs while attempting to retrieve the expiration time
+        :raises UnknownRecordError: If the OTP was not found
+        """
+
+    @abstractmethod
+    def reset(self):
+        """
+        Remove all `OTP`s from the repository
+
+        :raises RemovalError: If an error occurs while attempting to reset the repository
+        """

monkey/monkey_island/cc/repositories/mongo_otp_repository.py

@@ -0,0 +1,44 @@
+from pymongo import MongoClient
+
+from monkey_island.cc.server_utils.encryption import ILockableEncryptor
+
+from . import IOTPRepository
+from .consts import MONGO_OBJECT_ID_KEY
+from .errors import RemovalError, RetrievalError, StorageError, UnknownRecordError
+
+
+class MongoOTPRepository(IOTPRepository):
+    def __init__(
+        self,
+        mongo_client: MongoClient,
+        encryptor: ILockableEncryptor,
+    ):
+        self._encryptor = encryptor
+        self._otp_collection = mongo_client.monkey_island.otp
+        self._otp_collection.create_index("otp", unique=True)
+
+    def insert_otp(self, otp: str, expiration: float):
+        try:
+            encrypted_otp = self._encryptor.encrypt(otp.encode())
+            self._otp_collection.insert_one({"otp": encrypted_otp, "expiration_time": expiration})
+        except Exception as err:
+            raise StorageError(f"Error updating otp: {err}")
+
+    def get_expiration(self, otp: str) -> float:
+        try:
+            encrypted_otp = self._encryptor.encrypt(otp.encode())
+            otp_dict = self._otp_collection.find_one(
+                {"otp": encrypted_otp}, {MONGO_OBJECT_ID_KEY: False}
+            )
+        except Exception as err:
+            raise RetrievalError(f"Error retrieving otp: {err}")
+
+        if otp_dict is None:
+            raise UnknownRecordError("OTP not found")
+        return otp_dict["expiration_time"]
+
+    def reset(self):
+        try:
+            self._otp_collection.drop()
+        except Exception as err:
+            raise RemovalError(f"Error resetting the repository: {err}")

monkey/tests/unit_tests/monkey_island/cc/repositories/test_mongo_otp_repository.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+import mongomock
+import pytest
+
+from monkey_island.cc.repositories import (
+    IOTPRepository,
+    MongoOTPRepository,
+    RemovalError,
+    RetrievalError,
+    StorageError,
+    UnknownRecordError,
+)
+
+
+@dataclass
+class OTP:
+    otp: str
+    expiration_time: float
+
+
+OTPS = (
+    OTP(otp="test_otp_1", expiration_time=1),
+    OTP(otp="test_otp_2", expiration_time=2),
+    OTP(otp="test_otp_3", expiration_time=3),
+)
+
+
+@pytest.fixture
+def otp_repository(repository_encryptor) -> IOTPRepository:
+    return MongoOTPRepository(mongomock.MongoClient(), repository_encryptor)
+
+
+@pytest.fixture
+def error_raising_mongo_client() -> mongomock.MongoClient:
+    client = mongomock.MongoClient()
+    client.monkey_island = MagicMock(spec=mongomock.Database)
+    client.monkey_island.otp = MagicMock(spec=mongomock.Collection)
+    client.monkey_island.otp.insert_one = MagicMock(side_effect=Exception("insert failed"))
+    client.monkey_island.otp.find_one = MagicMock(side_effect=Exception("find failed"))
+    client.monkey_island.otp.delete_one = MagicMock(side_effect=Exception("delete failed"))
+    client.monkey_island.otp.drop = MagicMock(side_effect=Exception("drop failed"))
+
+    return client
+
+
+@pytest.fixture
+def error_raising_otp_repository(
+    error_raising_mongo_client, repository_encryptor
+) -> IOTPRepository:
+    return MongoOTPRepository(error_raising_mongo_client, repository_encryptor)
+
+
+def test_insert_otp(otp_repository: IOTPRepository):
+    otp_repository.insert_otp("test_otp", 1)
+    assert otp_repository.get_expiration("test_otp") == 1
+
+
+def test_insert_otp__prevents_duplicates(otp_repository: IOTPRepository):
+    otp_repository.insert_otp(OTPS[0].otp, OTPS[0].expiration_time)
+    with pytest.raises(StorageError):
+        otp_repository.insert_otp(OTPS[0].otp, OTPS[0].expiration_time)
+
+
+def test_insert_otp__prevents_duplicate_otp_with_differing_expiration(
+    otp_repository: IOTPRepository,
+):
+    otp_repository.insert_otp(OTPS[0].otp, 11)
+    with pytest.raises(StorageError):
+        otp_repository.insert_otp(OTPS[0].otp, 99)
+
+
+def test_insert_otp__raises_storage_error_if_error_occurs(
+    error_raising_otp_repository: IOTPRepository,
+):
+    with pytest.raises(StorageError):
+        error_raising_otp_repository.insert_otp("test_otp", 1)
+
+
+def test_get_expiration__raises_unknown_record_error_if_otp_does_not_exist(
+    otp_repository: IOTPRepository,
+):
+    with pytest.raises(UnknownRecordError):
+        otp_repository.get_expiration("test_otp")
+
+
+def test_get_expiration__returns_expiration_if_otp_exists(otp_repository: IOTPRepository):
+    otp_repository.insert_otp(OTPS[0].otp, OTPS[0].expiration_time)
+    assert otp_repository.get_expiration(OTPS[0].otp) == OTPS[0].expiration_time
+
+
+def test_get_expiration__raises_retrieval_error_if_error_occurs(
+    error_raising_otp_repository: IOTPRepository,
+):
+    with pytest.raises(RetrievalError):
+        error_raising_otp_repository.get_expiration("test_otp")
+
+
+def test_reset__deletes_all_otp(otp_repository: IOTPRepository):
+    for o in OTPS:
+        otp_repository.insert_otp(o.otp, o.expiration_time)
+
+    otp_repository.reset()
+
+    for o in OTPS:
+        with pytest.raises(UnknownRecordError):
+            otp_repository.get_expiration(o.otp)
+
+
+def test_reset__raises_removal_error_if_error_occurs(error_raising_otp_repository: IOTPRepository):
+    with pytest.raises(RemovalError):
+        error_raising_otp_repository.reset()

vulture_allowlist.py

@@ -18,7 +18,12 @@
 from infection_monkey.transport.http import FileServHTTPRequestHandler
 from monkey_island.cc.deployment import Deployment
 from monkey_island.cc.models import IslandMode, Machine
-from monkey_island.cc.repositories import IAgentEventRepository, MongoAgentEventRepository
+from monkey_island.cc.repositories import (
+    IAgentEventRepository,
+    IOTPRepository,
+    MongoAgentEventRepository,
+    MongoOTPRepository,
+)
 from monkey_island.cc.services.authentication_service.token import TokenValidator
 from monkey_island.cc.services.authentication_service.user import User
 from monkey_island.cc.services.reporting.exploitations.monkey_exploitation import MonkeyExploitation
@@ -149,3 +154,9 @@
 # Remove after #3137
 TokenValidator.validate_token
 _refresh_token_validator
+
+# Remove after #3078
+IOTPRepository.insert_otp
+IOTPRepository.get_expiration
+IOTPRepository.reset
+MongoOTPRepository