From 8feaf15939973142706087df93000efba993295e Mon Sep 17 00:00:00 2001 From: benyissa Date: Thu, 16 Nov 2023 17:12:16 +0100 Subject: [PATCH 1/3] implement logic with tests --- agent/metasploit_agent.py | 50 ++++++++++++++++++++++++ agent/utils.py | 69 ++++++++++++++++++++++++++++++++++ tests/metasploit_agent_test.py | 53 ++++++++++++++++++++++++++ 3 files changed, 172 insertions(+) diff --git a/agent/metasploit_agent.py b/agent/metasploit_agent.py index dc00312..6f5118d 100644 --- a/agent/metasploit_agent.py +++ b/agent/metasploit_agent.py @@ -1,6 +1,7 @@ """Ostorlab Agent implementation for metasploit""" import logging import socket +import ipaddress import time from typing import Any @@ -32,6 +33,7 @@ [*] Scanned 1 of 1 hosts (100% complete) [*] Auxiliary module execution completed """ +METASPLOIT_AGENT_KEY = b"agent_metasploit_asset" class Error(Exception): @@ -70,6 +72,11 @@ def process(self, message: m.Message) -> None: message: A message containing the path and the content of the file to be processed """ + + logger.info("processing message of selector : %s", message.selector) + if self._is_target_already_processed(message) is True: + return + client = utils.connect_msfrpc() cid = client.consoles.console().cid for entry in self._config: @@ -140,6 +147,49 @@ def process(self, message: m.Message) -> None: self._emit_results(module_instance, technical_detail) client.logout() + self._mark_target_as_processed(message) + logger.info("Done processing message of selector : %s", message.selector) + + def _is_target_already_processed(self, message: m.Message) -> bool: + """Checks if the target has already been processed before, relies on the redis server.""" + if message.data.get("url") is not None or message.data.get("name") is not None: + unicity_check_key = utils.get_unique_check_key(message) + if unicity_check_key is None: + return True + return self.set_is_member(key=METASPLOIT_AGENT_KEY, value=unicity_check_key) + + if message.data.get("host") is not None: + host = str(message.data.get("host")) + mask = message.data.get("mask") + if mask is not None: + addresses = ipaddress.ip_network(f"{host}/{mask}", strict=False) + return self.ip_network_exists( + key=METASPLOIT_AGENT_KEY, ip_range=addresses + ) + return self.set_is_member(key=METASPLOIT_AGENT_KEY, value=host) + logger.error("Unknown target %s", message) + return True + + def _mark_target_as_processed(self, message: m.Message) -> None: + """Mark the target as processed, relies on the redis server.""" + if message.data.get("url") is not None or message.data.get("name") is not None: + unicity_check_key = utils.get_unique_check_key(message) + if unicity_check_key is None: + return + + self.set_add(METASPLOIT_AGENT_KEY, unicity_check_key) + elif message.data.get("host") is not None: + host = str(message.data.get("host")) + mask = message.data.get("mask") + if mask is not None: + addresses = ipaddress.ip_network(f"{host}/{mask}", strict=False) + self.add_ip_network(key=METASPLOIT_AGENT_KEY, ip_range=addresses) + else: + self.set_add(METASPLOIT_AGENT_KEY, host) + else: + logger.error("Unknown target %s", message) + return + def _get_job_results( self, client: msfrpc.MsfRpcClient, job_uuid: int ) -> dict[str, Any] | list[str] | None: diff --git a/agent/utils.py b/agent/utils.py index 6b5cbda..26c1779 100644 --- a/agent/utils.py +++ b/agent/utils.py @@ -1,4 +1,6 @@ """Utilities for agent Metasploit""" +import dataclasses +from typing import cast from urllib import parse as urlparser import tenacity @@ -28,6 +30,14 @@ DEFAULT_PORT = 443 MSFRPCD_PWD = "Ostorlab123" PROCESS_TIMEOUT = 300 +DEFAULT_SCHEME = "https" + + +@dataclasses.dataclass +class Target: + host: str + scheme: str + port: int def _get_port(message: m.Message) -> int: @@ -38,6 +48,65 @@ def _get_port(message: m.Message) -> int: return DEFAULT_PORT +def _get_scheme(message: m.Message) -> str: + """Returns the schema to be used for the target.""" + protocol = message.data.get("protocol") + if protocol is not None: + return str(protocol) + + schema = message.data.get("schema") + if schema is None: + return DEFAULT_SCHEME + if schema in [ + "https?", + "ssl/https-alt?", + "ssl/https-alt", + "https-alt", + "https-alt?", + ]: + return "https" + if schema in ["http?", "http"]: + return "http" + return str(schema) + + +def get_unique_check_key(message: m.Message) -> str | None: + """Compute a unique key for a target""" + if message.data.get("url") is not None: + target = _get_target_from_url(message) + if target is not None: + return f"{target.scheme}_{target.host}_{target.port}" + elif message.data.get("name") is not None: + port = _get_port(message) + schema = _get_scheme(message) + domain = message.data["name"] + return f"{schema}_{domain}_{port}" + return None + + +def _get_target_from_url(message: m.Message) -> Target | None: + """Compute schema and port from a URL""" + url = message.data["url"] + parsed_url = urlparser.urlparse(url) + if parsed_url.scheme not in SCHEME_TO_PORT: + return None + schema = parsed_url.scheme or DEFAULT_SCHEME + schema = cast(str, schema) + domain_name = urlparser.urlparse(url).netloc + port = 0 + if len(parsed_url.netloc.split(":")) > 1: + domain_name = parsed_url.netloc.split(":")[0] + if ( + len(parsed_url.netloc.split(":")) > 0 + and parsed_url.netloc.split(":")[-1] != "" + ): + port = int(parsed_url.netloc.split(":")[-1]) + args_port = _get_port(message) + port = port or SCHEME_TO_PORT.get(schema) or args_port + target = Target(host=domain_name, scheme=schema, port=port) + return target + + def prepare_target(message: m.Message) -> tuple[str, int]: """Prepare targets based on type, if a domain name is provided, port and protocol are collected from the config.""" diff --git a/tests/metasploit_agent_test.py b/tests/metasploit_agent_test.py index 0f5827c..7cd2140 100644 --- a/tests/metasploit_agent_test.py +++ b/tests/metasploit_agent_test.py @@ -14,6 +14,7 @@ def testExploitCheck_whenSafe_returnNone( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit exploit check, case when target is safe""" @@ -30,6 +31,7 @@ def testExploitCheck_whenSafe_returnNone( def testAuxiliaryExecute_whenSafe_returnNone( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit auxiliary execute, case when target is safe""" @@ -45,6 +47,7 @@ def testAuxiliaryExecute_whenVulnerable_returnFindings( agent_instance: msf_agent.MetasploitAgent, mocker: plugin.MockerFixture, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit auxiliary execute, case when target is vulnerable""" @@ -82,6 +85,7 @@ def testExploitCheck_whenVulnerable_returnFindings( agent_instance: msf_agent.MetasploitAgent, mocker: plugin.MockerFixture, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], metasploitable_scan_message: message.Message, ) -> None: """Unit test for agent metasploit exploit check, case when target is vulnerable""" @@ -123,6 +127,7 @@ def testExploitCheck_whenVulnerable_returnConsoleOutput( agent_instance: msf_agent.MetasploitAgent, mocker: plugin.MockerFixture, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], metasploitable_scan_message: message.Message, ) -> None: """Unit test for agent metasploit exploit check, case when target is vulnerable (console output)""" @@ -164,6 +169,7 @@ def testExploitCheck_whenVulnerable_returnConsoleOutput( def testAuxiliaryPortScan_whenResultsFound_returnOpenPorts( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit auxiliary run, case when results are found""" @@ -185,6 +191,7 @@ def testAgent_whenMultipleModulesUsed_returnFindings( agent_multi_instance: msf_agent.MetasploitAgent, mocker: plugin.MockerFixture, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit auxiliary run, case when results are found""" @@ -221,6 +228,7 @@ def testAgent_whenMultipleModulesUsed_returnFindings( def testExploitCheck_whenCannotCheck_returnNone( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit exploit check, case when cannot check""" @@ -244,3 +252,48 @@ def testExploitCheck_whenDefaultAuxiliaryMessage_returnNone( agent_instance.process(scan_message) assert len(agent_mock) == 0 + + +@pytest.mark.parametrize( + "agent_instance", + [["exploit/windows/http/exchange_proxyshell_rce", []]], + indirect=True, +) +def testAgentNuclei_whenSameMessageSentTwice_shouldScanOnlyOnce( + agent_instance: msf_agent.MetasploitAgent, + agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], + scan_message: message.Message, + mocker: plugin.MockerFixture, +) -> None: + """Test MetasploitAgent agent should not scan the same message twice.""" + connect_msfrpc_mock = mocker.patch("agent.utils.connect_msfrpc") + + agent_instance.process(scan_message) + agent_instance.process(scan_message) + + connect_msfrpc_mock.assert_called_once() + + +@pytest.mark.parametrize( + "agent_instance", + [["exploit/windows/http/exchange_proxyshell_rce", []]], + indirect=True, +) +def testMetasploitAgent_whenUnknownTarget_shouldNotBeProcessed( + agent_instance: msf_agent.MetasploitAgent, + agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], + scan_message: message.Message, + mocker: plugin.MockerFixture, +) -> None: + """Test MetasploitAgent agent should not scan message with unknown target.""" + + connect_msfrpc_mock = mocker.patch("agent.utils.connect_msfrpc") + msg = message.Message.from_data( + "v3.asset.file", data={"path": "libagora-crypto.so"} + ) + + agent_instance.process(msg) + + connect_msfrpc_mock.assert_not_called() From 6da2029ca854fb298934b43c73d53f21cfbb4979 Mon Sep 17 00:00:00 2001 From: benyissa Date: Thu, 16 Nov 2023 17:19:06 +0100 Subject: [PATCH 2/3] change test name --- tests/metasploit_agent_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/metasploit_agent_test.py b/tests/metasploit_agent_test.py index 7cd2140..8476cdb 100644 --- a/tests/metasploit_agent_test.py +++ b/tests/metasploit_agent_test.py @@ -259,7 +259,7 @@ def testExploitCheck_whenDefaultAuxiliaryMessage_returnNone( [["exploit/windows/http/exchange_proxyshell_rce", []]], indirect=True, ) -def testAgentNuclei_whenSameMessageSentTwice_shouldScanOnlyOnce( +def testMetasploitAgent_whenSameMessageSentTwice_shouldScanOnlyOnce( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], agent_persist_mock: dict[str | bytes, str | bytes], From 64066074d842d1d17ca471cc81a55df4233950cc Mon Sep 17 00:00:00 2001 From: benyissa Date: Thu, 16 Nov 2023 17:28:15 +0100 Subject: [PATCH 3/3] fix failed test --- tests/metasploit_agent_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/metasploit_agent_test.py b/tests/metasploit_agent_test.py index 8476cdb..3dd1de9 100644 --- a/tests/metasploit_agent_test.py +++ b/tests/metasploit_agent_test.py @@ -245,6 +245,7 @@ def testExploitCheck_whenCannotCheck_returnNone( def testExploitCheck_whenDefaultAuxiliaryMessage_returnNone( agent_instance: msf_agent.MetasploitAgent, agent_mock: list[message.Message], + agent_persist_mock: dict[str | bytes, str | bytes], scan_message: message.Message, ) -> None: """Unit test for agent metasploit exploit check,