Feat: skip processed messages

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:

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."""

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"""
@@ -237,10 +245,56 @@ 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,
     case when console returns default auxiliary message"""
     agent_instance.process(scan_message)
 
     assert len(agent_mock) == 0
+
+
+@pytest.mark.parametrize(
+    "agent_instance",
+    [["exploit/windows/http/exchange_proxyshell_rce", []]],
+    indirect=True,
+)
+def testMetasploitAgent_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()