Merge branch '3122-test-island-logout' into develop

Issue #3122
PR #3175

envs/monkey_zoo/blackbox/gcp_test_machine_list.py

@@ -103,6 +103,7 @@
 SMB_PTH = {"europe-west3-a": ["mimikatz-15"]}
 
 GCP_SINGLE_TEST_LIST = {
+    "test_logout": {},
     "test_depth_2_a": DEPTH_2_A,
     "test_depth_1_a": DEPTH_1_A,
     "test_depth_3_a": DEPTH_3_A,

envs/monkey_zoo/blackbox/island_client/i_monkey_island_requests.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+
+
+class IMonkeyIslandRequests(ABC):
+    @abstractmethod
+    def get_token_from_server(self):
+        pass
+
+    @abstractmethod
+    def login(self):
+        pass
+
+    @abstractmethod
+    def get(self, url, data=None):
+        pass
+
+    @abstractmethod
+    def post(self, url, data):
+        pass
+
+    @abstractmethod
+    def put(self, url, data):
+        pass
+
+    @abstractmethod
+    def put_json(self, url, json: Dict):
+        pass
+
+    @abstractmethod
+    def post_json(self, url, json: Dict):
+        pass
+
+    @abstractmethod
+    def patch(self, url, data: Dict):
+        pass
+
+    @abstractmethod
+    def delete(self, url):
+        pass

envs/monkey_zoo/blackbox/island_client/monkey_island_client.py

@@ -6,7 +6,7 @@
 
 from common.credentials import Credentials
 from common.types import AgentID, MachineID
-from envs.monkey_zoo.blackbox.island_client.monkey_island_requests import MonkeyIslandRequests
+from envs.monkey_zoo.blackbox.island_client.i_monkey_island_requests import IMonkeyIslandRequests
 from envs.monkey_zoo.blackbox.test_configurations.test_configuration import TestConfiguration
 from monkey_island.cc.models import Agent, Machine, TerminateAllAgents
 
@@ -16,6 +16,7 @@
 ISLAND_LOG_ENDPOINT = "api/island/log"
 GET_MACHINES_ENDPOINT = "api/machines"
 GET_AGENT_EVENTS_ENDPOINT = "api/agent-events"
+LOGOUT_ENDPOINT = "api/logout"
 
 LOGGER = logging.getLogger(__name__)
 
@@ -26,8 +27,8 @@ def avoid_race_condition(func):
 
 
 class MonkeyIslandClient(object):
-    def __init__(self, server_address):
-        self.requests = MonkeyIslandRequests(server_address)
+    def __init__(self, requests: IMonkeyIslandRequests):
+        self.requests = requests
 
     def get_api_status(self):
         return self.requests.get("api")
@@ -177,3 +178,18 @@ def get_agent_events(self):
     def is_all_monkeys_dead(self):
         agents = self.get_agents()
         return all((a.stop_time is not None for a in agents))
+
+    def login(self):
+        try:
+            self.requests.login()
+            LOGGER.info("Logged into the Island.")
+        except Exception:
+            LOGGER.error("Failed to log into the Island.")
+            assert False
+
+    def logout(self):
+        if self.requests.post(LOGOUT_ENDPOINT, data=None).ok:
+            LOGGER.info("Logged out of the Island.")
+        else:
+            LOGGER.error("Failed to log out of the Island.")
+            assert False

envs/monkey_zoo/blackbox/island_client/monkey_island_requests.py

@@ -1,10 +1,10 @@
-import functools
 import logging
-from http import HTTPStatus
 from typing import Dict
 
 import requests
 
+from .i_monkey_island_requests import IMonkeyIslandRequests
+
 ISLAND_USERNAME = "test"
 ISLAND_PASSWORD = "testtest"
 LOGGER = logging.getLogger(__name__)
@@ -14,17 +14,20 @@ class InvalidRequestError(Exception):
     pass
 
 
-class MonkeyIslandRequests:
+class MonkeyIslandRequests(IMonkeyIslandRequests):
     def __init__(self, server_address):
         self.addr = f"https://{server_address}/"
-        self.token = self.try_get_token_from_server()
+        self.token = self._try_get_token_from_server()
 
-    def try_get_token_from_server(self):
+    def _try_get_token_from_server(self):
         try:
-            return self.try_set_island_to_credentials()
+            return self._try_set_island_to_credentials()
         except InvalidRequestError:
             return self.get_token_from_server()
 
+    def login(self):
+        self.token = self.get_token_from_server()
+
     def get_token_from_server(self):
         resp = requests.post(  # noqa: DUO123
             self.addr + "api/login",
@@ -38,7 +41,7 @@ def get_token_from_server(self):
         token = resp.json()["response"]["user"]["authentication_token"]
         return token
 
-    def try_set_island_to_credentials(self):
+    def _try_set_island_to_credentials(self):
         resp = requests.post(  # noqa: DUO123
             self.addr + "api/register",
             json={"username": ISLAND_USERNAME, "password": ISLAND_PASSWORD},
@@ -55,22 +58,6 @@ def try_set_island_to_credentials(self):
         token = resp.json()["response"]["user"]["authentication_token"]
         return token
 
-    class _Decorators:
-        @classmethod
-        def refresh_auth_token(cls, request_function):
-            @functools.wraps(request_function)
-            def request_function_wrapper(self, *args, **kwargs):
-                # noinspection PyArgumentList
-                resp = request_function(self, *args, **kwargs)
-                if resp.status_code == HTTPStatus.UNAUTHORIZED:
-                    self.token = self.get_token_from_server()
-                    resp = request_function(self, *args, **kwargs)
-
-                return resp
-
-            return request_function_wrapper
-
-    @_Decorators.refresh_auth_token
     def get(self, url, data=None):
         return requests.get(  # noqa: DUO123
             self.addr + url,
@@ -79,37 +66,31 @@ def get(self, url, data=None):
             verify=False,
         )
 
-    @_Decorators.refresh_auth_token
     def post(self, url, data):
         return requests.post(  # noqa: DUO123
             self.addr + url, data=data, headers=self.get_auth_header(), verify=False
         )
 
-    @_Decorators.refresh_auth_token
     def put(self, url, data):
         return requests.put(  # noqa: DUO123
             self.addr + url, data=data, headers=self.get_auth_header(), verify=False
         )
 
-    @_Decorators.refresh_auth_token
     def put_json(self, url, json: Dict):
         return requests.put(  # noqa: DUO123
             self.addr + url, json=json, headers=self.get_auth_header(), verify=False
         )
 
-    @_Decorators.refresh_auth_token
     def post_json(self, url, json: Dict):
         return requests.post(  # noqa: DUO123
             self.addr + url, json=json, headers=self.get_auth_header(), verify=False
         )
 
-    @_Decorators.refresh_auth_token
     def patch(self, url, data: Dict):
         return requests.patch(  # noqa: DUO123
             self.addr + url, data=data, headers=self.get_auth_header(), verify=False
         )
 
-    @_Decorators.refresh_auth_token
     def delete(self, url):
         return requests.delete(  # noqa: DUO123
             self.addr + url, headers=self.get_auth_header(), verify=False

envs/monkey_zoo/blackbox/island_client/reauthorizing_monkey_island_requests.py

@@ -0,0 +1,59 @@
+import functools
+from http import HTTPStatus
+from typing import Dict
+
+from .i_monkey_island_requests import IMonkeyIslandRequests
+
+
+class ReauthorizingMonkeyIslandRequests(IMonkeyIslandRequests):
+    def __init__(self, monkey_island_requests: IMonkeyIslandRequests):
+        self.requests = monkey_island_requests
+
+    def get_token_from_server(self):
+        return self.requests.get_token_from_server()
+
+    def login(self):
+        self.requests.login()
+
+    class _Decorators:
+        @classmethod
+        def refresh_auth_token(cls, request_function):
+            @functools.wraps(request_function)
+            def request_function_wrapper(self, *args, **kwargs):
+                # noinspection PyArgumentList
+                resp = request_function(self, *args, **kwargs)
+                if resp.status_code == HTTPStatus.UNAUTHORIZED:
+                    self.requests.login()
+                    resp = request_function(self, *args, **kwargs)
+
+                return resp
+
+            return request_function_wrapper
+
+    @_Decorators.refresh_auth_token
+    def get(self, url, data=None):
+        return self.requests.get(url, data=data)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def post(self, url, data):
+        return self.requests.post(url, data=data)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def put(self, url, data):
+        return self.requests.put(url, data=data)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def put_json(self, url, json: Dict):
+        return self.requests.put_json(url, json=json)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def post_json(self, url, json: Dict):
+        return self.requests.post_json(url, json=json)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def patch(self, url, data: Dict):
+        return self.requests.patch(url, data=data)  # noqa: DUO123
+
+    @_Decorators.refresh_auth_token
+    def delete(self, url):
+        return self.requests.delete(url)  # noqa: DUO123

envs/monkey_zoo/blackbox/test_blackbox.py

@@ -1,12 +1,24 @@
 import logging
 import os
+from http import HTTPStatus
 from time import sleep
 
 import pytest
 
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.analyzers.zerologon_analyzer import ZerologonAnalyzer
-from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
+from envs.monkey_zoo.blackbox.island_client.i_monkey_island_requests import IMonkeyIslandRequests
+from envs.monkey_zoo.blackbox.island_client.monkey_island_client import (
+    GET_AGENTS_ENDPOINT,
+    GET_MACHINES_ENDPOINT,
+    ISLAND_LOG_ENDPOINT,
+    LOGOUT_ENDPOINT,
+    MonkeyIslandClient,
+)
+from envs.monkey_zoo.blackbox.island_client.monkey_island_requests import MonkeyIslandRequests
+from envs.monkey_zoo.blackbox.island_client.reauthorizing_monkey_island_requests import (
+    ReauthorizingMonkeyIslandRequests,
+)
 from envs.monkey_zoo.blackbox.island_client.test_configuration_parser import get_target_ips
 from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
 from envs.monkey_zoo.blackbox.test_configurations import (
@@ -63,11 +75,17 @@ def wait_machine_bootup():
     sleep(MACHINE_BOOTUP_WAIT_SECONDS)
 
 
+@pytest.fixture
+def monkey_island_requests(island) -> IMonkeyIslandRequests:
+    return MonkeyIslandRequests(island)
+
+
 @pytest.fixture(scope="class")
-def island_client(island):
+def island_client(monkey_island_requests):
     client_established = False
     try:
-        island_client_object = MonkeyIslandClient(island)
+        requests = ReauthorizingMonkeyIslandRequests(monkey_island_requests)
+        island_client_object = MonkeyIslandClient(requests)
         client_established = island_client_object.get_api_status()
     except Exception:
         logging.exception("Got an exception while trying to establish connection to the Island.")
@@ -77,6 +95,34 @@ def island_client(island):
     yield island_client_object
 
 
+@pytest.mark.parametrize(
+    "authenticated_endpoint",
+    [
+        GET_AGENTS_ENDPOINT,
+        ISLAND_LOG_ENDPOINT,
+        GET_MACHINES_ENDPOINT,
+    ],
+)
+def test_logout(monkey_island_requests, authenticated_endpoint):
+    # Prove that we can't access authenticated endpoints without logging in
+    resp = monkey_island_requests.get(authenticated_endpoint)
+    assert resp.status_code == HTTPStatus.UNAUTHORIZED
+
+    # Prove that we can access authenticated endpoints after logging in
+    monkey_island_requests.login()
+    resp = monkey_island_requests.get(authenticated_endpoint)
+    assert resp.ok
+
+    # Log out - NOTE: This is an "out-of-band" call to logout. DO NOT call
+    # `monkey_island_request.logout()`. This could allow implementation details of the
+    # MonkeyIslandRequests class to cause false positives.
+    monkey_island_requests.post(LOGOUT_ENDPOINT, data=None)
+
+    # Prove that we can't access authenticated endpoints after logging out
+    resp = monkey_island_requests.get(authenticated_endpoint)
+    assert resp.status_code == HTTPStatus.UNAUTHORIZED
+
+
 # NOTE: These test methods are ordered to give time for the slower zoo machines
 # to boot up and finish starting services.
 @pytest.mark.usefixtures("island_client")