3078 use otp secret type

monkey/common/types/__init__.py

@@ -5,3 +5,4 @@
 from .networking import NetworkService, NetworkPort, PortStatus, SocketAddress, NetworkProtocol
 from .plugin_types import PluginName
 from .plugin_types import PluginVersion
+from .otp import OTP

monkey/common/types/otp.py

@@ -1,5 +1,5 @@
 from typing import TypeAlias
 
-from common.utils.i_secret_variable import ISecretVariable
+from pydantic import SecretStr
 
-OTP: TypeAlias = ISecretVariable
+OTP: TypeAlias = SecretStr

monkey/infection_monkey/island_api_client/http_island_api_client.py

@@ -16,8 +16,7 @@
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
 from common.common_consts.token_keys import ACCESS_TOKEN_KEY_NAME, REFRESH_TOKEN_KEY_NAME
 from common.credentials import Credentials
-from common.types import AgentID, JSONSerializable
-from common.types.otp import OTP
+from common.types import OTP, AgentID, JSONSerializable
 
 from . import IIslandAPIClient, IslandAPIRequestError
 from .http_client import HTTPClient

monkey/infection_monkey/monkey.py

@@ -34,12 +34,10 @@
 from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue, QueuedAgentEventPublisher
 from common.network.network_utils import get_my_ip_addresses, get_network_interfaces
 from common.tags.attack import T1082_ATTACK_TECHNIQUE_TAG
-from common.types import NetworkPort, SocketAddress
-from common.types.otp import OTP
+from common.types import OTP, NetworkPort, SocketAddress
 from common.utils.argparse_types import positive_int
 from common.utils.code_utils import del_key, secure_generate_random_string
 from common.utils.file_utils import create_secure_directory
-from common.utils.secret_variable import SecretVariable
 from infection_monkey.agent_event_handlers import (
     AgentEventForwarder,
     add_stolen_credentials_to_propagation_credentials_repository,
@@ -177,10 +175,10 @@ def _get_arguments(args):
     def _get_otp() -> OTP:
         # No need for a constant, this is a feature flag that will be removed.
         if OTP_FLAG not in os.environ:
-            return SecretVariable("PLACEHOLDER_OTP")
+            return OTP("PLACEHOLDER_OTP")
 
         try:
-            otp = SecretVariable(os.environ[AGENT_OTP_ENVIRONMENT_VARIABLE])
+            otp = OTP(os.environ[AGENT_OTP_ENVIRONMENT_VARIABLE])
         except KeyError:
             raise Exception(
                 f"Couldn't find {AGENT_OTP_ENVIRONMENT_VARIABLE} environmental variable. "

monkey/monkey_island/cc/services/authentication_service/__init__.py

@@ -1,4 +1,3 @@
-from .types import OTP
 from .account_role import AccountRole
 from .flask_resources import register_resources
 from .i_otp_generator import IOTPGenerator

monkey/monkey_island/cc/services/authentication_service/authentication_facade.py

@@ -5,6 +5,7 @@
 
 from flask_security import UserDatastore
 
+from common.types import OTP
 from common.utils.code_utils import secure_generate_random_string
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.models import IslandMode
@@ -15,7 +16,7 @@
 from . import AccountRole
 from .i_otp_repository import IOTPRepository
 from .token_parser import ParsedToken, TokenParser
-from .types import OTP, Token
+from .types import Token
 from .user import User
 
 OTP_EXPIRATION_TIME = 2 * 60  # 2 minutes
@@ -95,7 +96,7 @@ def generate_otp(self) -> OTP:
 
         The generated OTP is saved to the `IOTPRepository`
         """
-        otp = secure_generate_random_string(32, string.ascii_letters + string.digits + "._-")
+        otp = OTP(secure_generate_random_string(32, string.ascii_letters + string.digits + "._-"))
         expiration_time = time.monotonic() + OTP_EXPIRATION_TIME
         self._otp_repository.insert_otp(otp, expiration_time)
 

monkey/monkey_island/cc/services/authentication_service/authentication_service_otp_generator.py

@@ -1,6 +1,7 @@
+from common.types import OTP
+
 from .authentication_facade import AuthenticationFacade
 from .i_otp_generator import IOTPGenerator
-from .types import OTP
 
 
 class AuthenticationServiceOTPGenerator(IOTPGenerator):

monkey/monkey_island/cc/services/authentication_service/flask_resources/agent_otp.py

@@ -51,6 +51,6 @@ def get(self):
             raise RuntimeError("limiter has not been initialized")
         try:
             with AgentOTP.limiter:
-                return make_response({"otp": self._otp_generator.generate_otp()})
+                return make_response({"otp": self._otp_generator.generate_otp().get_secret_value()})
         except RateLimitExceeded:
             return make_response("Rate limit exceeded", HTTPStatus.TOO_MANY_REQUESTS)

monkey/monkey_island/cc/services/authentication_service/flask_resources/agent_otp_login.py

@@ -7,13 +7,12 @@
 from flask_security.registerable import register_user
 
 from common.common_consts.token_keys import ACCESS_TOKEN_KEY_NAME, REFRESH_TOKEN_KEY_NAME
-from common.types import AgentID
+from common.types import OTP, AgentID
 from common.utils.code_utils import secure_generate_random_string
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.services.authentication_service import AccountRole
 
 from ..authentication_facade import AuthenticationFacade
-from ..types import OTP
 
 
 class ArgumentParsingException(Exception):

monkey/monkey_island/cc/services/authentication_service/i_otp_generator.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from .types import OTP
+from common.types import OTP
 
 
 class IOTPGenerator(ABC):

monkey/monkey_island/cc/services/authentication_service/i_otp_repository.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from .types import OTP
+from common.types import OTP
 
 
 class IOTPRepository(ABC):

monkey/monkey_island/cc/services/authentication_service/mongo_otp_repository.py

@@ -4,6 +4,7 @@
 from bson.objectid import ObjectId
 from pymongo import MongoClient
 
+from common.types import OTP
 from monkey_island.cc.repositories import (
     MONGO_OBJECT_ID_KEY,
     RemovalError,
@@ -14,7 +15,6 @@
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor
 
 from .i_otp_repository import IOTPRepository
-from .types import OTP
 
 
 class MongoOTPRepository(IOTPRepository):
@@ -29,7 +29,7 @@ def __init__(
 
     def insert_otp(self, otp: OTP, expiration: float):
         try:
-            encrypted_otp = self._encryptor.encrypt(otp.encode())
+            encrypted_otp = self._encryptor.encrypt(otp.get_secret_value().encode())
             self._otp_collection.insert_one(
                 {"otp": encrypted_otp, "expiration_time": expiration, "used": False}
             )
@@ -71,8 +71,10 @@ def _get_otp_document(self, otp: OTP) -> Mapping[str, Any]:
 
     @lru_cache
     def _get_otp_object_id(self, otp: OTP) -> ObjectId:
+        otp_str = otp.get_secret_value()
+
         try:
-            encrypted_otp = self._encryptor.encrypt(otp.encode())
+            encrypted_otp = self._encryptor.encrypt(otp_str.encode())
             otp_dict = self._otp_collection.find_one({"otp": encrypted_otp}, [MONGO_OBJECT_ID_KEY])
         except Exception as err:
             raise RetrievalError(f"Error retrieving OTP: {err}")

monkey/monkey_island/cc/services/authentication_service/types.py

@@ -1,4 +1,3 @@
 from typing import TypeAlias
 
 Token: TypeAlias = str
-OTP: TypeAlias = str

monkey/monkey_island/cc/services/run_local_monkey.py

@@ -9,9 +9,9 @@
 from typing import Sequence
 
 from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
+from common.types import OTP
 from monkey_island.cc.repositories import IAgentBinaryRepository, RetrievalError
 from monkey_island.cc.server_utils.consts import ISLAND_PORT
-from monkey_island.cc.services.authentication_service import OTP
 
 logger = logging.getLogger(__name__)
 
@@ -73,7 +73,7 @@ def run_local_monkey(self, otp: OTP):
             port = ISLAND_PORT
 
             process_env = os.environ.copy()
-            process_env[AGENT_OTP_ENVIRONMENT_VARIABLE] = otp
+            process_env[AGENT_OTP_ENVIRONMENT_VARIABLE] = otp.get_secret_value()
             args = [str(dest_path), "m0nk3y", "-s", f"{ip}:{port}"]
             subprocess.Popen(args, cwd=self._data_dir, env=process_env)
         except Exception as exc:

monkey/tests/data_for_tests/otp.py

@@ -1,3 +1,3 @@
-from common.utils.secret_variable import SecretVariable
+from common.types import OTP
 
-OTP = SecretVariable("fake_otp")
+TEST_OTP = OTP("test_otp")

monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py

@@ -7,7 +7,7 @@
 import pytest
 import requests
 from tests.common.example_agent_configuration import AGENT_CONFIGURATION
-from tests.data_for_tests.otp import OTP
+from tests.data_for_tests.otp import TEST_OTP
 from tests.data_for_tests.propagation_credentials import CREDENTIALS_DICTS
 from tests.unit_tests.common.agent_plugins.test_agent_plugin_manifest import (
     FAKE_AGENT_MANIFEST_DICT,
@@ -107,7 +107,7 @@ def test_login__connection_error():
     api_client = build_api_client(http_client_stub)
 
     with pytest.raises(IslandAPIError):
-        api_client.login(OTP)
+        api_client.login(TEST_OTP)
 
 
 def test_login():
@@ -127,7 +127,7 @@ def test_login():
     }
     api_client = build_api_client(http_client_stub)
 
-    api_client.login(OTP)
+    api_client.login(TEST_OTP)
 
     assert http_client_stub.additional_headers[HTTPIslandAPIClient.TOKEN_HEADER_KEY] == auth_token
 
@@ -139,7 +139,7 @@ def test_login__bad_response():
     api_client = build_api_client(http_client_stub)
 
     with pytest.raises(IslandAPIAuthenticationError):
-        api_client.login(OTP)
+        api_client.login(TEST_OTP)
 
 
 def test_login__does_not_overwrite_additional_headers():
@@ -159,7 +159,7 @@ def test_login__does_not_overwrite_additional_headers():
     }
     api_client = build_api_client(http_client_stub)
 
-    api_client.login(OTP)
+    api_client.login(TEST_OTP)
 
     assert http_client_stub.additional_headers == {
         "Some-Header": "some value",

monkey/tests/unit_tests/infection_monkey/test_monkey.py

@@ -1,7 +1,7 @@
 import os
 
 import pytest
-from tests.data_for_tests.otp import OTP
+from tests.data_for_tests.otp import TEST_OTP
 
 from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
 from infection_monkey.model import OTP_FLAG
@@ -10,12 +10,12 @@
 
 @pytest.fixture(autouse=True)
 def configure_environment_variables(monkeypatch):
-    monkeypatch.setenv(AGENT_OTP_ENVIRONMENT_VARIABLE, OTP.get_secret_value())
+    monkeypatch.setenv(AGENT_OTP_ENVIRONMENT_VARIABLE, TEST_OTP.get_secret_value())
     monkeypatch.setenv(OTP_FLAG, True)
 
 
 def test_get_otp(monkeypatch):
-    assert InfectionMonkey._get_otp().get_secret_value() == OTP.get_secret_value()
+    assert InfectionMonkey._get_otp().get_secret_value() == TEST_OTP.get_secret_value()
     assert AGENT_OTP_ENVIRONMENT_VARIABLE not in os.environ
 
 

monkey/tests/unit_tests/monkey_island/cc/services/authentication_service/flask_resources/test_agent_otp.py

@@ -2,12 +2,11 @@
 from typing import Callable
 
 import pytest
+from tests.data_for_tests.otp import TEST_OTP
 
 from monkey_island.cc.services.authentication_service import IOTPGenerator
 from monkey_island.cc.services.authentication_service.flask_resources.agent_otp import AgentOTP
 
-OTP = "supersecretpassword"
-
 
 @pytest.fixture
 def make_otp_request(flask_client):
@@ -20,8 +19,8 @@ def _make_otp_request():
 
 
 def test_agent_otp__successful(make_otp_request: Callable, mock_otp_generator: IOTPGenerator):
-    mock_otp_generator.generate_otp.return_value = OTP
+    mock_otp_generator.generate_otp.return_value = TEST_OTP
     response = make_otp_request()
 
     assert response.status_code == HTTPStatus.OK
-    assert response.json["otp"] == OTP
+    assert response.json["otp"] == TEST_OTP.get_secret_value()

monkey/tests/unit_tests/monkey_island/cc/services/authentication_service/flask_resources/test_agent_otp_login.py

@@ -2,6 +2,7 @@
 from uuid import UUID
 
 import pytest
+from tests.data_for_tests.otp import TEST_OTP
 from tests.unit_tests.monkey_island.conftest import get_url_for_resource
 
 from common.common_consts.token_keys import ACCESS_TOKEN_KEY_NAME, REFRESH_TOKEN_KEY_NAME
@@ -23,7 +24,7 @@ def _agent_otp_login(request_body):
 
 
 def test_agent_otp_login__successful(agent_otp_login):
-    response = agent_otp_login({"agent_id": AGENT_ID, "otp": "supersecretpassword"})
+    response = agent_otp_login({"agent_id": AGENT_ID, "otp": TEST_OTP.get_secret_value()})
 
     assert response.status_code == HTTPStatus.OK
     assert ACCESS_TOKEN_KEY_NAME in response.json["response"]["user"]
@@ -37,8 +38,8 @@ def test_agent_otp_login__successful(agent_otp_login):
         [],
         {"otp": ""},
         {"agent_id": AGENT_ID},
-        {"agent_id": "", "otp": "supersecretpassword"},
-        {"agent_id": "1234", "otp": "supersecretpassword"},
+        {"agent_id": "", "otp": TEST_OTP.get_secret_value()},
+        {"agent_id": "1234", "otp": TEST_OTP.get_secret_value()},
     ],
 )
 def test_invalid_request(agent_otp_login, data):

monkey/tests/unit_tests/monkey_island/cc/services/authentication_service/test_authentication_service.py

@@ -8,6 +8,7 @@
 from flask_security import UserDatastore
 from tests.common import StubDIContainer
 
+from common.types import OTP
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.models import IslandMode
 from monkey_island.cc.repositories import UnknownRecordError
@@ -235,7 +236,7 @@ def test_authorize_otp(
     get_expiration_return_value: int,
     otp_is_valid_expected_value: bool,
 ):
-    otp = "secret"
+    otp = OTP("secret")
 
     freezer.move_to(TIME)
 
@@ -250,7 +251,7 @@ def test_authorize_otp__unknown_otp(
     authentication_facade: AuthenticationFacade,
     mock_otp_repository: IOTPRepository,
 ):
-    otp = "secret"
+    otp = OTP("secret")
 
     mock_otp_repository.otp_is_used.side_effect = UnknownRecordError(f"Unknown otp {otp}")
     mock_otp_repository.set_used.side_effect = UnknownRecordError(f"Unknown otp {otp}")