Merge branch '3078-rate-limit' into develop

Issue #3078
PR #3189

Author: mssalvatore

envs/monkey_zoo/blackbox/test_blackbox.py

@@ -1,9 +1,11 @@
 import logging
 import os
 from http import HTTPStatus
+from threading import Thread
 from time import sleep
 
 import pytest
+import requests
 
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.analyzers.zerologon_analyzer import ZerologonAnalyzer
@@ -38,6 +40,9 @@
     start_machines,
     stop_machines,
 )
+from monkey_island.cc.services.authentication_service.flask_resources.agent_otp import (
+    MAX_OTP_REQUESTS_PER_SECOND,
+)
 
 DEFAULT_TIMEOUT_SECONDS = 2 * 60 + 30
 MACHINE_BOOTUP_WAIT_SECONDS = 30
@@ -48,7 +53,11 @@
 
 @pytest.fixture(autouse=True, scope="session")
 def GCPHandler(request, no_gcp, gcp_machines_to_start):
-    if not no_gcp:
+    if no_gcp:
+        return
+    if len(gcp_machines_to_start) == 0:
+        LOGGER.info("No GCP machines to start.")
+    else:
         LOGGER.info(f"MACHINES TO START: {gcp_machines_to_start}")
 
         try:
@@ -156,6 +165,27 @@ def test_logout_invalidates_all_tokens(island):
     assert resp.status_code == HTTPStatus.UNAUTHORIZED
 
 
+def test_agent_otp_rate_limit(island):
+    threads = []
+    response_codes = []
+    agent_otp_endpoint = f"https://{island}/api/agent-otp"
+
+    def make_request():
+        response = requests.get(agent_otp_endpoint, verify=False)  # noqa: DUO123
+        response_codes.append(response.status_code)
+
+    for _ in range(0, MAX_OTP_REQUESTS_PER_SECOND + 1):
+        t = Thread(target=make_request, daemon=True)
+        t.start()
+        threads.append(t)
+
+    for t in threads:
+        t.join()
+
+    assert response_codes.count(HTTPStatus.OK) == MAX_OTP_REQUESTS_PER_SECOND
+    assert response_codes.count(HTTPStatus.TOO_MANY_REQUESTS) == 1
+
+
 # NOTE: These test methods are ordered to give time for the slower zoo machines
 # to boot up and finish starting services.
 # noinspection PyUnresolvedReferences

monkey/infection_monkey/exploit/i_agent_otp_provider.py

@@ -14,5 +14,6 @@ def get_otp(self) -> str:
         Get a one-time password (OTP)
 
         :return: An OTP
+        :raises RuntimeError: If an OTP cannot be retrieved
         """
         pass

monkey/infection_monkey/exploit/island_api_agent_otp_provider.py

@@ -1,4 +1,6 @@
-from infection_monkey.island_api_client import IIslandAPIClient
+from time import sleep
+
+from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIRequestLimitExceededError
 
 from .i_agent_otp_provider import IAgentOTPProvider
 
@@ -8,4 +10,14 @@ def __init__(self, island_api_client: IIslandAPIClient):
         self._island_api_client = island_api_client
 
     def get_otp(self) -> str:
-        return self._island_api_client.get_otp()
+        # Note: Using too large of a delay here will cause the exploiter
+        # threads/processes to hang
+        for _ in range(6):
+            try:
+                return self._island_api_client.get_otp()
+            except IslandAPIRequestLimitExceededError:
+                sleep(0.5)
+                continue
+            except Exception as err:
+                raise RuntimeError(err)
+        raise RuntimeError("Failed to get OTP: Too many requests.")

monkey/infection_monkey/island_api_client/__init__.py

@@ -4,6 +4,8 @@
     IslandAPIRequestError,
     IslandAPIRequestFailedError,
     IslandAPITimeoutError,
+    IslandAPIAuthenticationError,
+    IslandAPIRequestLimitExceededError,
 )
 from .i_island_api_client import IIslandAPIClient
 from .abstract_island_api_client_factory import AbstractIslandAPIClientFactory

monkey/infection_monkey/island_api_client/http_island_api_client.py

@@ -1,6 +1,7 @@
 import functools
 import json
 import logging
+from http import HTTPStatus
 from pprint import pformat
 from typing import Any, Dict, List, Sequence
 
@@ -20,7 +21,11 @@
 
 from . import IIslandAPIClient, IslandAPIRequestError
 from .http_client import HTTPClient
-from .island_api_client_errors import IslandAPIAuthenticationError, IslandAPIResponseParsingError
+from .island_api_client_errors import (
+    IslandAPIAuthenticationError,
+    IslandAPIRequestLimitExceededError,
+    IslandAPIResponseParsingError,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -116,6 +121,8 @@ def get_agent_binary(self, operating_system: OperatingSystem) -> bytes:
     @handle_authentication_token_expiration
     def get_otp(self) -> str:
         response = self._http_client.get("/agent-otp")
+        if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            raise IslandAPIRequestLimitExceededError("Too many requests to get OTP.")
         return response.json()["otp"]
 
     @handle_response_parsing_errors

monkey/infection_monkey/island_api_client/i_island_api_client.py

@@ -61,6 +61,7 @@ def get_otp(self) -> str:
                                        Island due to an issue in the request sent from the client
         :raises IslandAPIRequestFailedError: If an error occurs while attempting to connect to the
                                              Island due to an error on the server
+        :raises IslandAPIRequestLimitExceededError: If the request limit for OTPs has been exceeded
         :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the Island
         :raises IslandAPIError: If an unexpected error occurs while attempting to get an OTP
         :return: The OTP

monkey/infection_monkey/island_api_client/island_api_client_errors.py

@@ -50,3 +50,13 @@ class IslandAPIResponseParsingError(IslandAPIError):
     """
     Raised when IslandAPIClient fails to parse the response
     """
+
+    pass
+
+
+class IslandAPIRequestLimitExceededError(IslandAPIError):
+    """
+    Raised when the API request fails due to rate limiting
+    """
+
+    pass

monkey/monkey_island/Pipfile

@@ -17,6 +17,7 @@ requests = ">=2.24"
 ring = ">=0.7.3"
 Flask-RESTful = ">=0.3.8"
 Flask = ">=1.1"
+Flask-Limiter = "*"
 Werkzeug = ">=1.0.1"
 pyaescrypt = "*"
 python-dateutil = "*"