Merge pull request #1818 from guardicore/1604-custom-pba

1604 custom pba
This commit is contained in:
Mike Salvatore 2022-03-30 09:40:44 -04:00 committed by GitHub
commit 581ece577d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 151 additions and 142 deletions

View File

@ -89,15 +89,6 @@ class Configuration(object):
keep_tunnel_open_time = 60
###########################
# post breach actions
###########################
post_breach_actions = []
custom_PBA_linux_cmd = ""
custom_PBA_windows_cmd = ""
PBA_linux_filename = None
PBA_windows_filename = None
###########################
# testing configuration
###########################

View File

@ -1,7 +1,7 @@
import logging
import threading
import time
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple
from infection_monkey.credential_store import ICredentialsStore
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
@ -16,6 +16,7 @@ from infection_monkey.utils.threading import create_daemon_thread, interruptible
from infection_monkey.utils.timer import Timer
from . import Exploiter, IPScanner, Propagator
from .option_parsing import custom_pba_is_enabled
CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5
CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5
@ -154,9 +155,9 @@ class AutomatedMaster(IMaster):
),
)
pba_thread = create_daemon_thread(
target=self._run_plugins,
target=self._run_pbas,
name="PBAThread",
args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba),
args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]),
)
credential_collector_thread.start()
@ -196,10 +197,6 @@ class AutomatedMaster(IMaster):
name = pba[0]
options = pba[1]
# TEMPORARY; TO AVOID ERRORS SINCE THIS ISN'T IMPLEMENTED YET
if name == "Custom":
return
for pba_data in self._puppet.run_pba(name, options):
self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
@ -212,6 +209,14 @@ class AutomatedMaster(IMaster):
self._puppet.run_payload(name, options, self._stop)
def _run_pbas(
self, plugins: Iterable[Any], callback: Callable[[Any], None], custom_pba_options: Mapping
):
self._run_plugins(plugins, "post-breach action", callback)
if custom_pba_is_enabled(custom_pba_options):
self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback)
def _run_plugins(
self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None]
):

View File

@ -0,0 +1,13 @@
from typing import Dict
from infection_monkey.utils.environment import is_windows_os
def custom_pba_is_enabled(pba_options: Dict) -> bool:
if not is_windows_os():
if pba_options["linux_command"]:
return True
else:
if pba_options["windows_command"]:
return True
return False

View File

@ -51,6 +51,7 @@ from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs
from infection_monkey.post_breach.actions.timestomping import Timestomping
from infection_monkey.post_breach.actions.use_signed_scripts import SignedScriptProxyExecution
from infection_monkey.post_breach.actions.use_trap_command import TrapCommand
from infection_monkey.post_breach.custom_pba import CustomPBA
from infection_monkey.puppet.puppet import Puppet
from infection_monkey.system_singleton import SystemSingleton
from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
@ -315,6 +316,9 @@ class InfectionMonkey:
ClearCommandHistory(self._telemetry_messenger),
PluginType.POST_BREACH_ACTION,
)
puppet.load_plugin(
"CustomPBA", CustomPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION
)
puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD)

View File

@ -0,0 +1 @@
from .custom_pba import CustomPBA

View File

@ -1,10 +1,11 @@
import logging
import os
from typing import Dict, Iterable
from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION
from common.utils.attack_utils import ScanStatus
from infection_monkey.config import WormConfiguration
from infection_monkey.control import ControlClient
from infection_monkey.i_puppet import PostBreachData
from infection_monkey.network.tools import get_interface_to_target
from infection_monkey.post_breach.pba import PBA
from infection_monkey.telemetry.attack.t1105_telem import T1105Telem
@ -19,52 +20,50 @@ DIR_CHANGE_WINDOWS = "cd %s & "
DIR_CHANGE_LINUX = "cd %s ; "
class UsersPBA(PBA):
class CustomPBA(PBA):
"""
Defines user's configured post breach action.
"""
def __init__(self, telemetry_messenger: ITelemetryMessenger):
super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
super(CustomPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION)
self.filename = ""
if not is_windows_os():
# Add linux commands to PBA's
if WormConfiguration.PBA_linux_filename:
self.filename = WormConfiguration.PBA_linux_filename
if WormConfiguration.custom_PBA_linux_cmd:
# Add change dir command, because user will try to access his file
self.command = (
DIR_CHANGE_LINUX % get_monkey_dir_path()
) + WormConfiguration.custom_PBA_linux_cmd
elif WormConfiguration.custom_PBA_linux_cmd:
self.command = WormConfiguration.custom_PBA_linux_cmd
else:
def run(self, options: Dict) -> Iterable[PostBreachData]:
self._set_options(options)
return super().run(options)
def _set_options(self, options: Dict):
# Required for attack telemetry
self.current_server = options["current_server"]
if is_windows_os():
# Add windows commands to PBA's
if WormConfiguration.PBA_windows_filename:
self.filename = WormConfiguration.PBA_windows_filename
if WormConfiguration.custom_PBA_windows_cmd:
if options["windows_filename"]:
self.filename = options["windows_filename"]
if options["windows_command"]:
# Add change dir command, because user will try to access his file
self.command = (
DIR_CHANGE_WINDOWS % get_monkey_dir_path()
) + WormConfiguration.custom_PBA_windows_cmd
elif WormConfiguration.custom_PBA_windows_cmd:
self.command = WormConfiguration.custom_PBA_windows_cmd
self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + options[
"windows_command"
]
elif options["windows_command"]:
self.command = options["windows_command"]
else:
# Add linux commands to PBA's
if options["linux_filename"]:
self.filename = options["linux_filename"]
if options["linux_command"]:
# Add change dir command, because user will try to access his file
self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + options[
"linux_command"
]
elif options["linux_command"]:
self.command = options["linux_command"]
def _execute_default(self):
if self.filename:
UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename)
return super(UsersPBA, self)._execute_default()
@staticmethod
def should_run(class_name):
if not is_windows_os():
if WormConfiguration.PBA_linux_filename or WormConfiguration.custom_PBA_linux_cmd:
return True
else:
if WormConfiguration.PBA_windows_filename or WormConfiguration.custom_PBA_windows_cmd:
return True
return False
self.download_pba_file(get_monkey_dir_path(), self.filename)
return super(CustomPBA, self)._execute_default()
def download_pba_file(self, dst_dir, filename):
"""
@ -84,11 +83,11 @@ class UsersPBA(PBA):
if not status:
status = ScanStatus.USED
self._telemetry_messenger.send_telemetry(
self.telemetry_messenger.send_telemetry(
T1105Telem(
status,
WormConfiguration.current_server.split(":")[0],
get_interface_to_target(WormConfiguration.current_server.split(":")[0]),
self.current_server.split(":")[0],
get_interface_to_target(self.current_server.split(":")[0]),
filename,
)
)

View File

@ -452,11 +452,13 @@ class ConfigService:
for pba in config.get("post_breach_actions", []):
formatted_pbas_config[pba] = {}
formatted_pbas_config["Custom"] = {
config["custom_pbas"] = {
"linux_command": config.get(flat_linux_command_field, ""),
"linux_filename": config.get(flat_linux_filename_field, ""),
"windows_command": config.get(flat_windows_command_field, ""),
"windows_filename": config.get(flat_windows_filename_field, ""),
# Current server is used for attack telemetry
"current_server": config.get("current_server"),
}
config["post_breach_actions"] = formatted_pbas_config

View File

@ -2,19 +2,20 @@ from unittest.mock import MagicMock
import pytest
from infection_monkey.post_breach.actions.users_custom_pba import UsersPBA
from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA
MONKEY_DIR_PATH = "/dir/to/monkey/"
CUSTOM_LINUX_CMD = "command-for-linux"
CUSTOM_LINUX_FILENAME = "filename-for-linux"
CUSTOM_WINDOWS_CMD = "command-for-windows"
CUSTOM_WINDOWS_FILENAME = "filename-for-windows"
CUSTOM_SERVER = "10.10.10.10:5000"
@pytest.fixture
@pytest.fixture(autouse=True)
def fake_monkey_dir_path(monkeypatch):
monkeypatch.setattr(
"infection_monkey.post_breach.actions.users_custom_pba.get_monkey_dir_path",
"infection_monkey.post_breach.custom_pba.custom_pba.get_monkey_dir_path",
lambda: MONKEY_DIR_PATH,
)
@ -22,7 +23,7 @@ def fake_monkey_dir_path(monkeypatch):
@pytest.fixture
def set_os_linux(monkeypatch):
monkeypatch.setattr(
"infection_monkey.post_breach.actions.users_custom_pba.is_windows_os",
"infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os",
lambda: False,
)
@ -30,106 +31,92 @@ def set_os_linux(monkeypatch):
@pytest.fixture
def set_os_windows(monkeypatch):
monkeypatch.setattr(
"infection_monkey.post_breach.actions.users_custom_pba.is_windows_os",
"infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os",
lambda: True,
)
@pytest.fixture
def mock_UsersPBA_linux_custom_file_and_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
CUSTOM_LINUX_CMD,
)
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.PBA_linux_filename",
CUSTOM_LINUX_FILENAME,
)
return UsersPBA(MagicMock())
def fake_custom_pba_linux_options():
return {
"linux_command": CUSTOM_LINUX_CMD,
"linux_filename": CUSTOM_LINUX_FILENAME,
"windows_command": "",
"windows_filename": "",
# Current server is used for attack telemetry
"current_server": CUSTOM_SERVER,
}
def test_command_linux_custom_file_and_cmd(
mock_UsersPBA_linux_custom_file_and_cmd,
):
def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux):
pba = CustomPBA(MagicMock())
pba._set_options(fake_custom_pba_linux_options)
expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}"
assert mock_UsersPBA_linux_custom_file_and_cmd.command == expected_command
assert pba.command == expected_command
assert pba.filename == CUSTOM_LINUX_FILENAME
@pytest.fixture
def mock_UsersPBA_windows_custom_file_and_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd",
CUSTOM_WINDOWS_CMD,
)
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.PBA_windows_filename",
CUSTOM_WINDOWS_FILENAME,
)
return UsersPBA(MagicMock())
def fake_custom_pba_windows_options():
return {
"linux_command": "",
"linux_filename": "",
"windows_command": CUSTOM_WINDOWS_CMD,
"windows_filename": CUSTOM_WINDOWS_FILENAME,
# Current server is used for attack telemetry
"current_server": CUSTOM_SERVER,
}
def test_command_windows_custom_file_and_cmd(
mock_UsersPBA_windows_custom_file_and_cmd,
):
def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows):
pba = CustomPBA(MagicMock())
pba._set_options(fake_custom_pba_windows_options)
expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}"
assert mock_UsersPBA_windows_custom_file_and_cmd.command == expected_command
assert pba.command == expected_command
assert pba.filename == CUSTOM_WINDOWS_FILENAME
@pytest.fixture
def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None)
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.PBA_linux_filename",
CUSTOM_LINUX_FILENAME,
)
return UsersPBA(MagicMock())
def fake_options_files_only():
return {
"linux_command": "",
"linux_filename": CUSTOM_LINUX_FILENAME,
"windows_command": "",
"windows_filename": CUSTOM_WINDOWS_FILENAME,
# Current server is used for attack telemetry
"current_server": CUSTOM_SERVER,
}
def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file):
expected_command = ""
assert mock_UsersPBA_linux_custom_file.command == expected_command
@pytest.mark.parametrize("os", [set_os_linux, set_os_windows])
def test_files_only(fake_options_files_only, os):
pba = CustomPBA(MagicMock())
pba._set_options(fake_options_files_only)
assert pba.command == ""
@pytest.fixture
def mock_UsersPBA_windows_custom_file(set_os_windows, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None)
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.PBA_windows_filename",
CUSTOM_WINDOWS_FILENAME,
)
return UsersPBA(MagicMock())
def fake_options_commands_only():
return {
"linux_command": CUSTOM_LINUX_CMD,
"linux_filename": "",
"windows_command": CUSTOM_WINDOWS_CMD,
"windows_filename": "",
# Current server is used for attack telemetry
"current_server": CUSTOM_SERVER,
}
def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file):
expected_command = ""
assert mock_UsersPBA_windows_custom_file.command == expected_command
def test_commands_only(fake_options_commands_only, set_os_linux):
pba = CustomPBA(MagicMock())
pba._set_options(fake_options_commands_only)
assert pba.command == CUSTOM_LINUX_CMD
assert pba.filename == ""
@pytest.fixture
def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
CUSTOM_LINUX_CMD,
)
monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_linux_filename", None)
return UsersPBA(MagicMock())
def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd):
expected_command = CUSTOM_LINUX_CMD
assert mock_UsersPBA_linux_custom_cmd.command == expected_command
@pytest.fixture
def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch):
monkeypatch.setattr(
"infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd",
CUSTOM_WINDOWS_CMD,
)
monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_windows_filename", None)
return UsersPBA(MagicMock())
def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd):
expected_command = CUSTOM_WINDOWS_CMD
assert mock_UsersPBA_windows_custom_cmd.command == expected_command
def test_commands_only_windows(fake_options_commands_only, set_os_windows):
pba = CustomPBA(MagicMock())
pba._set_options(fake_options_commands_only)
assert pba.command == CUSTOM_WINDOWS_CMD
assert pba.filename == ""

View File

@ -66,12 +66,6 @@ def test_format_config_for_agent__pbas(flat_monkey_config):
"ScheduleJobs": {},
"Timestomping": {},
"AccountDiscovery": {},
"Custom": {
"linux_command": "bash test.sh",
"windows_command": "powershell test.ps1",
"linux_filename": "test.sh",
"windows_filename": "test.ps1",
},
}
ConfigService.format_flat_config_for_agent(flat_monkey_config)
@ -84,6 +78,19 @@ def test_format_config_for_agent__pbas(flat_monkey_config):
assert "PBA_windows_filename" not in flat_monkey_config
def test_format_config_for_custom_pbas(flat_monkey_config):
custom_config = {
"linux_command": "bash test.sh",
"windows_command": "powershell test.ps1",
"linux_filename": "test.sh",
"windows_filename": "test.ps1",
"current_server": "10.197.94.72:5000",
}
ConfigService.format_flat_config_for_agent(flat_monkey_config)
assert flat_monkey_config["custom_pbas"] == custom_config
def test_get_config_propagation_credentials_from_flat_config(flat_monkey_config):
expected_creds = {
"exploit_lm_hash_list": ["lm_hash_1", "lm_hash_2"],