Merge pull request #2034 from guardicore/1960-file-retrieval-error-todo

1960 file retrieval error todo

monkey/monkey_island/cc/repository/__init__.py

@@ -1,6 +1,6 @@
 from .errors import RetrievalError
-from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository
+from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository
-from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError
+from .i_agent_binary_repository import IAgentBinaryRepository
 from .agent_binary_repository import AgentBinaryRepository
 from .i_agent_configuration_repository import IAgentConfigurationRepository
 from .file_agent_configuration_repository import FileAgentConfigurationRepository

monkey/monkey_island/cc/repository/agent_binary_repository.py

@@ -1,6 +1,6 @@
 from typing import BinaryIO
 
-from . import AgentRetrievalError, FileRetrievalError, IAgentBinaryRepository, IFileRepository
+from . import IAgentBinaryRepository, IFileRepository, RetrievalError
 
 LINUX_AGENT_FILE_NAME = "monkey-linux-64"
 WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe"
@@ -20,8 +20,8 @@ class AgentBinaryRepository(IAgentBinaryRepository):
         try:
             agent_binary = self._file_repository.open_file(filename)
             return agent_binary
-        except FileRetrievalError as err:
+        except Exception as err:
-            raise AgentRetrievalError(
+            raise RetrievalError(
                 f"An error occurred while retrieving the {filename}"
                 f" agent binary from {self._file_repository}: {err}"
             )

monkey/monkey_island/cc/repository/file_agent_configuration_repository.py

@@ -5,6 +5,7 @@ from common.configuration import (
     AgentConfiguration,
     AgentConfigurationSchema,
 )
+from monkey_island.cc import repository
 from monkey_island.cc.repository import (
     IAgentConfigurationRepository,
     IFileRepository,
@@ -25,10 +26,10 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository):
                 configuration_json = f.read().decode()
 
             return self._schema.loads(configuration_json)
-        # TODO: Handle FileRetrievalError vs FileNotFoundError
+        except repository.FileNotFoundError:
-        #       https://github.com/guardicore/monkey/blob/e8001d8cf76340e42bf17ff62523bd2d85fc4841/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py#L47-L50
-        except RetrievalError:
             return self._schema.loads(DEFAULT_AGENT_CONFIGURATION)
+        except Exception as err:
+            raise RetrievalError(f"Error retrieving the agent configuration: {err}")
 
     def store_configuration(self, agent_configuration: AgentConfiguration):
         configuration_json = self._schema.dumps(agent_configuration)

monkey/monkey_island/cc/repository/file_storage/__init__.py

@@ -1,2 +1,2 @@
-from .i_file_repository import IFileRepository, FileRetrievalError
+from .i_file_repository import IFileRepository, FileNotFoundError
 from .local_storage_file_repository import LocalStorageFileRepository

monkey/monkey_island/cc/repository/file_storage/i_file_repository.py

@@ -4,7 +4,7 @@ from typing import BinaryIO
 from monkey_island.cc.repository import RetrievalError
 
 
-class FileRetrievalError(RetrievalError):
+class FileNotFoundError(RetrievalError):
     pass
 
 
@@ -30,7 +30,8 @@ class IFileRepository(metaclass=abc.ABCMeta):
 
         :param unsafe_file_name: An unsanitized file name that identifies the file to be opened
         :return: A file-like object providing access to the file's contents
-        :raises FileRetrievalError: if the file cannot be opened
+        :raises FileNotFoundError: if the file does not exist
+        :raises RetrievalError: if the file exists but cannot be retrieved
         """
         pass
 

monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py

@@ -4,9 +4,10 @@ from pathlib import Path
 from typing import BinaryIO
 
 from common.utils.file_utils import get_all_regular_files_in_directory
+from monkey_island.cc.repository import RetrievalError
 from monkey_island.cc.server_utils.file_utils import create_secure_directory
 
-from . import FileRetrievalError, IFileRepository
+from . import IFileRepository, i_file_repository
 
 logger = logging.getLogger(__name__)
 
@@ -43,12 +44,15 @@ class LocalStorageFileRepository(IFileRepository):
         try:
             logger.debug(f"Opening {safe_file_path}")
             return open(safe_file_path, "rb")
-        except OSError as err:
+        except FileNotFoundError as err:
-            # TODO: The interface should make a destinction between file not found and an error when
+            # Wrap Python's FileNotFound error, which is-an OSError, in repository.FileNotFoundError
-            #       retrieving a file that should exist. The built-in `FileNotFoundError` is not
+            raise i_file_repository.FileNotFoundError(
-            #       sufficient because it inherits from `OSError` and the interface does not
+                f'The requested file "{unsafe_file_name}" does not exist: {err}'
-            #       guarantee that the file is stored on the local file system.
+            )
-            raise FileRetrievalError(f"Failed to retrieve file {safe_file_path}: {err}") from err
+        except Exception as err:
+            raise RetrievalError(
+                f'Error retrieving file "{unsafe_file_name}" from the repository: {err}'
+            )
 
     def delete_file(self, unsafe_file_name: str):
         safe_file_path = self._get_safe_file_path(unsafe_file_name)

monkey/monkey_island/cc/repository/i_agent_binary_repository.py

@@ -2,10 +2,6 @@ import abc
 from typing import BinaryIO
 
 
-class AgentRetrievalError(IOError):
-    pass
-
-
 class IAgentBinaryRepository(metaclass=abc.ABCMeta):
     """
     A repository that retrieves the agent binaries
@@ -17,6 +13,7 @@ class IAgentBinaryRepository(metaclass=abc.ABCMeta):
         Retrieve linux agent binary
 
         :return: A file-like object that represents the linux agent binary
+        :raises RetrievalError: If the agent binary could not be retrieved
         """
 
     @abc.abstractmethod
@@ -25,4 +22,5 @@ class IAgentBinaryRepository(metaclass=abc.ABCMeta):
         Retrieve windows agent binary
 
         :return: A file-like object that represents the windows agent binary
+        :raises RetrievalError: If the agent binary could not be retrieved
         """

monkey/monkey_island/cc/repository/i_agent_configuration_repository.py

@@ -14,8 +14,8 @@ class IAgentConfigurationRepository(ABC):
         Retrieve the agent configuration from the repository
 
         :return: The agent configuration as retrieved from the repository, or the default
-                 configuration if none could be retrieved.
+                 configuration if the repository is empty
-        :raises RetrievalError: if the configuration can not be retrieved
+        :raises RetrievalError: if the configuration could not be retrieved
         """
         pass
 

monkey/monkey_island/cc/resources/agent_binaries.py

@@ -2,7 +2,7 @@ import logging
 
 from flask import make_response, send_file
 
-from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository
+from monkey_island.cc.repository import IAgentBinaryRepository, RetrievalError
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 
 logger = logging.getLogger(__name__)
@@ -31,10 +31,10 @@ class AgentBinaries(AbstractResource):
             file = agent_binaries[os]()
 
             return send_file(file, mimetype="application/octet-stream")
-        except AgentRetrievalError as err:
-            logger.error(err)
-            return make_response({"error": str(err)}, 500)
         except KeyError as err:
             error_msg = f'No Agents are available for unsupported operating system "{os}": {err}'
             logger.error(error_msg)
             return make_response({"error": error_msg}, 404)
+        except RetrievalError as err:
+            logger.error(err)
+            return make_response({"error": str(err)}, 500)

monkey/monkey_island/cc/resources/pba_file_download.py

@@ -2,7 +2,8 @@ import logging
 
 from flask import make_response, send_file
 
-from monkey_island.cc.repository import FileRetrievalError, IFileRepository
+from monkey_island.cc import repository
+from monkey_island.cc.repository import IFileRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 
 logger = logging.getLogger(__file__)
@@ -24,7 +25,6 @@ class PBAFileDownload(AbstractResource):
 
             # `send_file()` handles the closing of the open file.
             return send_file(file, mimetype="application/octet-stream")
-        except FileRetrievalError as err:
+        except repository.FileNotFoundError as err:
-            error_msg = f"Failed to open file {filename}: {err}"
+            logger.error(str(err))
-            logger.error(error_msg)
+            return make_response({"error": str(err)}, 404)
-            return make_response({"error": error_msg}, 404)

monkey/monkey_island/cc/resources/pba_file_upload.py

@@ -5,7 +5,8 @@ from flask import Response, make_response, request, send_file
 from werkzeug.utils import secure_filename as sanitize_filename
 
 from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH
-from monkey_island.cc.repository import FileRetrievalError, IFileRepository
+from monkey_island.cc import repository
+from monkey_island.cc.repository import IFileRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
 from monkey_island.cc.services.config import ConfigService
@@ -57,10 +58,9 @@ class FileUpload(AbstractResource):
 
             # `send_file()` handles the closing of the open file.
             return send_file(file, mimetype="application/octet-stream")
-        except FileRetrievalError as err:
+        except repository.FileNotFoundError as err:
-            error_msg = f"Failed to open file {filename}: {err}"
+            logger.error(str(err))
-            logger.error(error_msg)
+            return make_response({"error": str(err)}, 404)
-            return make_response({"error": error_msg}, 404)
 
     @jwt_required
     def post(self, target_os):

monkey/monkey_island/cc/services/initialize.py

@@ -6,12 +6,12 @@ from common.aws import AWSInstance
 from common.utils.file_utils import get_binary_io_sha256_hash
 from monkey_island.cc.repository import (
     AgentBinaryRepository,
-    AgentRetrievalError,
     FileAgentConfigurationRepository,
     IAgentBinaryRepository,
     IAgentConfigurationRepository,
     IFileRepository,
     LocalStorageFileRepository,
+    RetrievalError,
 )
 from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.services import AWSService
@@ -76,7 +76,7 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository):
             agent_binary = get_agent_binary()
             binary_sha256_hash = get_binary_io_sha256_hash(agent_binary)
             agent_hashes[os] = binary_sha256_hash
-        except AgentRetrievalError as err:
+        except RetrievalError as err:
             logger.error(f"No agent available for {os}: {err}")
 
     for os, binary_sha256_hash in agent_hashes.items():

monkey/monkey_island/cc/services/run_local_monkey.py

@@ -5,7 +5,7 @@ import subprocess
 from pathlib import Path
 from shutil import copyfileobj
 
-from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository
+from monkey_island.cc.repository import IAgentBinaryRepository, RetrievalError
 from monkey_island.cc.server_utils.consts import ISLAND_PORT
 from monkey_island.cc.services.utils.network_utils import local_ip_addresses
 
@@ -29,7 +29,7 @@ class LocalMonkeyRunService:
             }
 
             agent_binary = agents[platform.system().lower()]()
-        except AgentRetrievalError as err:
+        except RetrievalError as err:
             logger.error(
                 f"No Agent can be retrieved for the specified operating system"
                 f'"{operating_system}"'

monkey/tests/monkey_island/__init__.py

@@ -1 +1,3 @@
 from .single_file_repository import SingleFileRepository
+from .mock_file_repository import MockFileRepository, FILE_CONTENTS, FILE_NAME
+from .open_error_file_repository import OpenErrorFileRepository

monkey/tests/monkey_island/mock_file_repository.py

@@ -0,0 +1,28 @@
+import io
+from typing import BinaryIO
+
+from monkey_island.cc import repository
+from monkey_island.cc.repository import IFileRepository
+
+FILE_NAME = "test_file"
+FILE_CONTENTS = b"HelloWorld!"
+
+
+class MockFileRepository(IFileRepository):
+    def __init__(self):
+        self._file = io.BytesIO(FILE_CONTENTS)
+
+    def save_file(self, unsafe_file_name: str, file_contents: BinaryIO):
+        pass
+
+    def open_file(self, unsafe_file_name: str) -> BinaryIO:
+        if unsafe_file_name != FILE_NAME:
+            raise repository.FileNotFoundError()
+
+        return self._file
+
+    def delete_file(self, unsafe_file_name: str):
+        pass
+
+    def delete_all_files(self):
+        pass

monkey/tests/monkey_island/open_error_file_repository.py

@@ -0,0 +1,10 @@
+from typing import BinaryIO
+
+from monkey_island.cc.repository import RetrievalError
+
+from . import MockFileRepository
+
+
+class OpenErrorFileRepository(MockFileRepository):
+    def open_file(self, unsafe_file_name: str) -> BinaryIO:
+        raise RetrievalError("Error retrieving file")

monkey/tests/monkey_island/single_file_repository.py

@@ -1,7 +1,8 @@
 import io
 from typing import BinaryIO
 
-from monkey_island.cc.repository import FileRetrievalError, IFileRepository
+from monkey_island.cc import repository
+from monkey_island.cc.repository import IFileRepository
 
 
 class SingleFileRepository(IFileRepository):
@@ -13,7 +14,7 @@ class SingleFileRepository(IFileRepository):
 
     def open_file(self, unsafe_file_name: str) -> BinaryIO:
         if self._file is None:
-            raise FileRetrievalError()
+            raise repository.FileNotFoundError()
         return self._file
 
     def delete_file(self, unsafe_file_name: str):

monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py

@@ -1,8 +1,9 @@
+import pytest
 from tests.common.example_agent_configuration import AGENT_CONFIGURATION
-from tests.monkey_island import SingleFileRepository
+from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository
 
 from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema
-from monkey_island.cc.repository import FileAgentConfigurationRepository
+from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError
 
 
 def test_store_agent_config():
@@ -24,3 +25,10 @@ def test_get_default_agent_config():
     retrieved_agent_configuration = repository.get_configuration()
 
     assert retrieved_agent_configuration == default_agent_configuration
+
+
+def test_get_agent_config_retrieval_error():
+    repository = FileAgentConfigurationRepository(OpenErrorFileRepository())
+
+    with pytest.raises(RetrievalError):
+        repository.get_configuration()

monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py

@@ -1,10 +1,12 @@
 import io
 from pathlib import Path
+from unittest.mock import Mock, patch
 
 import pytest
 from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions
 
-from monkey_island.cc.repository import FileRetrievalError, LocalStorageFileRepository
+from monkey_island.cc import repository
+from monkey_island.cc.repository import LocalStorageFileRepository
 from monkey_island.cc.server_utils.file_utils import is_windows_os
 
 
@@ -131,5 +133,13 @@ def test_remove_nonexistant_file(tmp_path):
 def test_open_nonexistant_file(tmp_path):
     fss = LocalStorageFileRepository(tmp_path)
 
-    with pytest.raises(FileRetrievalError):
+    with pytest.raises(repository.FileNotFoundError):
         fss.open_file("nonexistant_file.txt")
+
+
+def test_open_locked_file(tmp_path, monkeypatch):
+    fss = LocalStorageFileRepository(tmp_path)
+
+    with patch("builtins.open", Mock(side_effect=Exception())):
+        with pytest.raises(repository.RetrievalError):
+            fss.open_file("locked_file.txt")

monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py

@@ -2,11 +2,14 @@ from unittest.mock import MagicMock
 
 import flask_jwt_extended
 import pytest
+from tests.common import StubDIContainer
+from tests.monkey_island import OpenErrorFileRepository
 from tests.unit_tests.monkey_island.conftest import init_mock_app
 
 import monkey_island.cc.app
 import monkey_island.cc.resources.auth.auth
 import monkey_island.cc.resources.island_mode
+from monkey_island.cc.repository import IFileRepository
 
 
 @pytest.fixture
@@ -38,3 +41,12 @@ def get_mock_app(container):
     flask_jwt_extended.JWTManager(app)
 
     return app
+
+
+@pytest.fixture
+def open_error_flask_client(build_flask_client):
+    container = StubDIContainer()
+    container.register(IFileRepository, OpenErrorFileRepository)
+
+    with build_flask_client(container) as flask_client:
+        yield flask_client

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py

@@ -1,36 +1,11 @@
-import io
-from typing import BinaryIO
-
 import pytest
 from tests.common import StubDIContainer
+from tests.monkey_island import FILE_CONTENTS, FILE_NAME, MockFileRepository
 from tests.unit_tests.monkey_island.conftest import get_url_for_resource
 
-from monkey_island.cc.repository import FileRetrievalError, IFileRepository
+from monkey_island.cc.repository import IFileRepository
 from monkey_island.cc.resources.pba_file_download import PBAFileDownload
 
-FILE_NAME = "test_file"
-FILE_CONTENTS = b"HelloWorld!"
-
-
-class MockFileRepository(IFileRepository):
-    def __init__(self):
-        self._file = io.BytesIO(FILE_CONTENTS)
-
-    def save_file(self, unsafe_file_name: str, file_contents: BinaryIO):
-        pass
-
-    def open_file(self, unsafe_file_name: str) -> BinaryIO:
-        if unsafe_file_name != FILE_NAME:
-            raise FileRetrievalError()
-
-        return self._file
-
-    def delete_file(self, unsafe_file_name: str):
-        pass
-
-    def delete_all_files(self):
-        pass
-
 
 @pytest.fixture
 def flask_client(build_flask_client):
@@ -56,3 +31,11 @@ def test_file_download_endpoint_404(tmp_path, flask_client):
     resp = flask_client.get(download_url)
 
     assert resp.status_code == 404
+
+
+def test_file_download_endpoint_500(tmp_path, open_error_flask_client):
+    download_url = get_url_for_resource(PBAFileDownload, filename="test")
+
+    resp = open_error_flask_client.get(download_url)
+
+    assert resp.status_code == 500

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py

@@ -98,6 +98,15 @@ def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_conf
     assert resp.status_code == 404
 
 
+@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE])
+def test_file_download_endpoint_500(open_error_flask_client, pba_os):
+    url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py")
+
+    resp = open_error_flask_client.get(url)
+
+    assert resp.status_code == 500
+
+
 @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE])
 def test_pba_file_upload_endpoint(
     flask_client, pba_os, mock_get_config_value, mock_set_config_value