diff --git a/monkey/infection_monkey/model/__init__.py b/monkey/infection_monkey/model/__init__.py index 206f41f9db5..313fd062549 100644 --- a/monkey/infection_monkey/model/__init__.py +++ b/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" diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index f474ae86964..4a322ee5961 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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))}") diff --git a/monkey/tests/unit_tests/infection_monkey/test_monkey.py b/monkey/tests/unit_tests/infection_monkey/test_monkey.py new file mode 100644 index 00000000000..7f77c565245 --- /dev/null +++ b/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"