From f9936fe65dc488029957e9bfab52649bdd7eb2ab Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 16 Mar 2022 11:55:46 -0400 Subject: [PATCH] Agent: Add connect() method to IPowerShellClient --- monkey/infection_monkey/exploit/powershell.py | 6 ++- .../powershell_utils/powershell_client.py | 21 ++++++--- .../exploit/test_powershell.py | 46 +++++++++++++------ 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/monkey/infection_monkey/exploit/powershell.py b/monkey/infection_monkey/exploit/powershell.py index 41b9d9d00..d18a5c982 100644 --- a/monkey/infection_monkey/exploit/powershell.py +++ b/monkey/infection_monkey/exploit/powershell.py @@ -127,9 +127,10 @@ class PowerShellExploiter(HostExploiter): encryption=ENCRYPTION_AUTO, ssl=use_ssl, ) - # TODO: Report login attempt or find a better way of detecting if SSL is enabled - PowerShellClient(self.host.ip_addr, credentials, auth_options) + # TODO: Report login attempt or find a better way of detecting if SSL is enabled + client = PowerShellClient(self.host.ip_addr, credentials, auth_options) + client.connect() def _authenticate_via_brute_force( self, credentials: List[Credentials], auth_options: List[AuthOptions] @@ -137,6 +138,7 @@ class PowerShellExploiter(HostExploiter): for (creds, opts) in zip(credentials, auth_options): try: client = PowerShellClient(self.host.ip_addr, creds, opts) + client.connect() logger.info( f"Successfully logged into {self.host.ip_addr} using Powershell. User: " f"{creds.username}, Secret Type: {creds.secret_type.name}" diff --git a/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py b/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py index 80aefee00..c0ae8b260 100644 --- a/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py +++ b/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py @@ -54,6 +54,10 @@ def format_password(credentials: Credentials) -> Optional[str]: class IPowerShellClient(Protocol, metaclass=abc.ABCMeta): + @abc.abstractmethod + def connect(self) -> str: + pass + @abc.abstractmethod def execute_cmd(self, cmd: str) -> str: pass @@ -72,14 +76,19 @@ class PowerShellClient(IPowerShellClient): _set_sensitive_packages_log_level_to_error() self._ip_addr = ip_addr + self._credentials = credentials + self._auth_options = auth_options + self._client = None + + def connect(self): self._client = Client( - ip_addr, - username=credentials.username, - password=format_password(credentials), + self._ip_addr, + username=self._credentials.username, + password=format_password(self._credentials), cert_validation=False, - auth=auth_options.auth_type, - encryption=auth_options.encryption, - ssl=auth_options.ssl, + auth=self._auth_options.auth_type, + encryption=self._auth_options.encryption, + ssl=self._auth_options.ssl, connection_timeout=CONNECTION_TIMEOUT, ) diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py index d3d516353..f75c57f17 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -58,8 +58,11 @@ def powershell_exploiter(monkeypatch): def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_arguments): - mock_powershell_client = MagicMock(side_effect=Exception) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + mock_powershell_client = MagicMock() + mock_powershell_client.connect = MagicMock(side_effect=Exception) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) assert not exploit_result.exploitation_success @@ -74,8 +77,12 @@ def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments else: raise Exception - mock_powershell_client = MagicMock(side_effect=allow_http) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + mock_powershell_client = MagicMock() + mock_powershell_client.connect = MagicMock(side_effect=allow_http) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) + powershell_exploiter.exploit_host(**powershell_arguments) for call_args in mock_powershell_client.call_args_list: @@ -89,8 +96,12 @@ def test_powershell_https(monkeypatch, powershell_exploiter, powershell_argument else: raise Exception - mock_powershell_client = MagicMock(side_effect=allow_https) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + mock_powershell_client = MagicMock() + mock_powershell_client.connect = MagicMock(side_effect=allow_https) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) + powershell_exploiter.exploit_host(**powershell_arguments) for call_args in mock_powershell_client.call_args_list: @@ -99,8 +110,11 @@ def test_powershell_https(monkeypatch, powershell_exploiter, powershell_argument def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_arguments): - mock_powershell_client = MagicMock(side_effect=AuthenticationErrorForTests) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + mock_powershell_client = MagicMock() + mock_powershell_client.connect = MagicMock(side_effect=AuthenticationErrorForTests) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) assert not exploit_result.exploitation_success @@ -142,13 +156,14 @@ def test_failed_copy(monkeypatch, powershell_exploiter, powershell_arguments): def test_failed_monkey_execution(monkeypatch, powershell_exploiter, powershell_arguments): - mock_client = MagicMock() - mock_client.execute_cmd_as_detached_process = MagicMock( + mock_powershell_client = MagicMock() + mock_powershell_client.execute_cmd_as_detached_process = MagicMock( side_effect=Exception("EXECUTION FAILED") ) - mock_powershell_client = MagicMock(side_effect=authenticate(mock_client)) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) assert exploit_result.exploitation_success is True @@ -172,8 +187,11 @@ def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, po # login attempt doesn't throw an exception, signifying that login with credentials was # successful. connection_attempts = [True, Exception, Exception, Exception, Exception, Exception, True] - mock_client = MagicMock(side_effect=connection_attempts) - monkeypatch.setattr(powershell, "PowerShellClient", mock_client) + mock_powershell_client = MagicMock(side_effect=connection_attempts) + mock_powershell_client.connect = MagicMock(side_effect=connection_attempts) + monkeypatch.setattr( + powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) + ) exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)