Merge branch '3077-get-otp-from-env' into develop

Issue #3077
PR #3121

monkey/infection_monkey/model/__init__.py

@@ -1,3 +1,6 @@
+# TODO: Remove this before closing #3077
+OTP_FLAG = "AGENT_OTP_FLAG"
+
 MONKEY_ARG = "m0nk3y"
 DROPPER_ARG = "dr0pp3r"
 ID_STRING = "M0NK3Y3XPL0ITABLE"

monkey/infection_monkey/monkey.py

@@ -30,12 +30,13 @@
 )
 from common.agent_plugins import AgentPluginType
 from common.agent_registration_data import AgentRegistrationData
+from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
 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.utils.argparse_types import positive_int
-from common.utils.code_utils import secure_generate_random_string
+from common.utils.code_utils import del_key, secure_generate_random_string
 from common.utils.file_utils import create_secure_directory
 from infection_monkey.agent_event_handlers import (
     AgentEventForwarder,
@@ -95,6 +96,7 @@
 from infection_monkey.utils.signal_handler import register_signal_handlers, reset_signal_handlers
 
 from .heart import Heart
+from .model import OTP_FLAG
 from .plugin_event_forwarder import PluginEventForwarder
 
 logger = logging.getLogger(__name__)
@@ -114,8 +116,7 @@ def __init__(self, args, ipc_logger_queue: multiprocessing.Queue, log_path: Path
         self._manager = context.Manager()
 
         self._opts = self._get_arguments(args)
-        # TODO read the otp from an env variable
-        self._otp = "hard-coded-otp"
+        self._otp = self._get_otp()
 
         self._ipc_logger_queue = ipc_logger_queue
 
@@ -169,6 +170,26 @@ def _get_arguments(args):
 
         return opts
 
+    @staticmethod
+    def _get_otp():
+        # No need for a constant, this is a feature flag that will be removed.
+        if OTP_FLAG not in os.environ:
+            return "PLACEHOLDER_OTP"
+
+        try:
+            otp = os.environ[AGENT_OTP_ENVIRONMENT_VARIABLE]
+        except KeyError:
+            raise Exception(
+                f"Couldn't find {AGENT_OTP_ENVIRONMENT_VARIABLE} environmental variable."
+                f"Without an OTP the agent will fail to authenticate!"
+            )
+
+        # SECURITY: There's no need to leave this floating around in a place as visible as
+        # environment variables for any longer than necessary.
+        del_key(os.environ, AGENT_OTP_ENVIRONMENT_VARIABLE)
+
+        return otp
+
     # TODO: By the time we finish 2292, _connect_to_island_api() may not need to return `server`
     def _connect_to_island_api(self) -> Tuple[SocketAddress, IIslandAPIClient]:
         logger.debug(f"Trying to wake up with servers: {', '.join(map(str, self._opts.servers))}")

monkey/tests/unit_tests/infection_monkey/test_monkey.py

@@ -0,0 +1,35 @@
+import os
+
+import pytest
+from tests.data_for_tests.otp import OTP
+
+from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
+from infection_monkey.model import OTP_FLAG
+from infection_monkey.monkey import InfectionMonkey
+
+
+@pytest.fixture(autouse=True)
+def configure_environment_variables(monkeypatch):
+    monkeypatch.setenv(AGENT_OTP_ENVIRONMENT_VARIABLE, OTP)
+    monkeypatch.setenv(OTP_FLAG, True)
+
+
+def test_get_otp(monkeypatch):
+    assert InfectionMonkey._get_otp() == OTP
+    assert AGENT_OTP_ENVIRONMENT_VARIABLE not in os.environ
+
+
+def test_get_otp__no_otp(monkeypatch):
+    monkeypatch.delenv(AGENT_OTP_ENVIRONMENT_VARIABLE)
+    with pytest.raises(Exception):
+        InfectionMonkey._get_otp()
+
+
+def test_get_otp__feature_flag_disabled(monkeypatch):
+    try:
+        monkeypatch.delenv(OTP_FLAG)
+    except KeyError:
+        pass
+
+    # No need for a constant, this code is testing a feature flag that will be removed.
+    assert InfectionMonkey._get_otp() == "PLACEHOLDER_OTP"