Merge pull request #2114 from guardicore/2004-validate-custom-pba-config

Custom PBA configuration validation

monkey/common/agent_configuration/agent_sub_configuration_schemas.py

@@ -1,4 +1,6 @@
-from marshmallow import Schema, fields, post_load
+import re
+
+from marshmallow import Schema, ValidationError, fields, post_load, validate, validates
 
 from .agent_sub_configurations import (
     CustomPBAConfiguration,
@@ -13,12 +15,48 @@ from .agent_sub_configurations import (
 )
 from .utils import freeze_lists
 
+valid_windows_custom_pba_filename_regex = re.compile(r"^[^<>:\"\\\/|?*]*[^<>:\"\\\/|?* \.]+$|^$")
+valid_linux_custom_pba_filename_regex = re.compile(r"^[^\0/]*$")
+
 
 class CustomPBAConfigurationSchema(Schema):
     linux_command = fields.Str()
-    linux_filename = fields.Str()
+    linux_filename = fields.Str(
+        validate=validate.Regexp(regex=valid_linux_custom_pba_filename_regex)
+    )
     windows_command = fields.Str()
-    windows_filename = fields.Str()
+    windows_filename = fields.Str(
+        validate=validate.Regexp(regex=valid_windows_custom_pba_filename_regex)
+    )
+
+    @validates("windows_filename")
+    def validate_windows_filename_not_reserved(self, windows_filename):
+        # filename shouldn't start with any of these and be followed by a period
+        if windows_filename.split(".")[0].upper() in [
+            "CON",
+            "PRN",
+            "AUX",
+            "NUL",
+            "COM1",
+            "COM2",
+            "COM3",
+            "COM4",
+            "COM5",
+            "COM6",
+            "COM7",
+            "COM8",
+            "COM9",
+            "LPT1",
+            "LPT2",
+            "LPT3",
+            "LPT4",
+            "LPT5",
+            "LPT6",
+            "LPT7",
+            "LPT8",
+            "LPT9",
+        ]:
+            raise ValidationError("Invalid Windows filename: reserved name used")
 
     @post_load
     def _make_custom_pba_configuration(self, data, **kwargs):

monkey/common/agent_configuration/agent_sub_configurations.py

@@ -4,6 +4,21 @@ from typing import Dict, Tuple
 
 @dataclass(frozen=True)
 class CustomPBAConfiguration:
+    """
+    A configuration for custom post-breach actions
+
+    Attributes:
+        :param linux_command: Command to run on Linux victim machines. If a file is uploaded,
+                              use this field to change its permissions, execute it, and/or delete it
+                              Example: `chmod +x file.sh; ./file.sh; rm file.sh`
+        :param linux_filename: Name of the file to upload on Linux victim machines
+        :param windows_command: Command to run on Windows victim machines. If a file is uploaded,
+                                use this field to change its permissions, execute it, and/or delete
+                                it
+                                Example: `file.bat & del file.bat`
+        :param windows_filename: Name of the file to upload on Windows victim machines
+    """
+
     linux_command: str
     linux_filename: str
     windows_command: str

monkey/tests/unit_tests/common/configuration/test_agent_configuration.py

@@ -2,6 +2,7 @@ import json
 from copy import deepcopy
 
 import pytest
+from marshmallow import ValidationError
 from tests.common.example_agent_configuration import (
     AGENT_CONFIGURATION,
     BLOCKED_IPS,
@@ -68,6 +69,44 @@ def test_custom_pba_configuration_schema():
     assert config.windows_filename == WINDOWS_FILENAME
 
 
+def test_custom_pba_configuration_schema__empty_filenames_allowed():
+    schema = CustomPBAConfigurationSchema()
+
+    empty_filename_configuration = CUSTOM_PBA_CONFIGURATION.copy()
+    empty_filename_configuration.update({"linux_filename": "", "windows_filename": ""})
+
+    config = schema.load(empty_filename_configuration)
+
+    assert config.linux_command == LINUX_COMMAND
+    assert config.linux_filename == ""
+    assert config.windows_command == WINDOWS_COMMAND
+    assert config.windows_filename == ""
+
+
+@pytest.mark.parametrize("linux_filename", ["/", "/abc/", "\0"])
+def test_custom_pba_configuration_schema__invalid_linux_filename(linux_filename):
+    schema = CustomPBAConfigurationSchema()
+
+    invalid_filename_configuration = CUSTOM_PBA_CONFIGURATION.copy()
+    invalid_filename_configuration["linux_filename"] = linux_filename
+
+    with pytest.raises(ValidationError):
+        schema.load(invalid_filename_configuration)
+
+
+@pytest.mark.parametrize(
+    "windows_filename", ["CON", "CON.txt", "con.abc.pdf", " ", "abc.", "a?b", "d\\e"]
+)
+def test_custom_pba_configuration_schema__invalid_windows_filename(windows_filename):
+    schema = CustomPBAConfigurationSchema()
+
+    invalid_filename_configuration = CUSTOM_PBA_CONFIGURATION.copy()
+    invalid_filename_configuration["windows_filename"] = windows_filename
+
+    with pytest.raises(ValidationError):
+        schema.load(invalid_filename_configuration)
+
+
 def test_scan_target_configuration():
     schema = ScanTargetConfigurationSchema()
 

vulture_allowlist.py

@@ -251,3 +251,4 @@ IFindingRepository.get_findings
 key_list
 simulation
 netmap
+validate_windows_filename_not_reserved