Merge remote-tracking branch 'origin/dev' into feature/2646_Update-pre-commit-to-check-for-valid-copyright

src/primaite/VERSION

@@ -1 +1 @@
-3.0.0b9
+3.0.0

src/primaite/game/agent/actions.py

@@ -11,7 +11,7 @@ AbstractAction. The ActionManager is responsible for:
 """
 import itertools
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
 from gymnasium import spaces
 
@@ -871,6 +871,74 @@ class NetworkPortDisableAction(AbstractAction):
         return ["network", "node", target_nodename, "network_interface", port_id, "disable"]
 
 
+class NodeNMAPPingScanAction(AbstractAction):
+    """Action which performs an NMAP ping scan."""
+
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
+        super().__init__(manager=manager)
+
+    def form_request(self, source_node: str, target_ip_address: Union[str, List[str]]) -> List[str]:  # noqa
+        return [
+            "network",
+            "node",
+            source_node,
+            "application",
+            "NMAP",
+            "ping_scan",
+            {"target_ip_address": target_ip_address},
+        ]
+
+
+class NodeNMAPPortScanAction(AbstractAction):
+    """Action which performs an NMAP port scan."""
+
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
+        super().__init__(manager=manager)
+
+    def form_request(
+        self,
+        source_node: str,
+        target_ip_address: Union[str, List[str]],
+        target_protocol: Optional[Union[str, List[str]]] = None,
+        target_port: Optional[Union[str, List[str]]] = None,
+    ) -> List[str]:  # noqa
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
+        return [
+            "network",
+            "node",
+            source_node,
+            "application",
+            "NMAP",
+            "port_scan",
+            {"target_ip_address": target_ip_address, "target_port": target_port, "target_protocol": target_protocol},
+        ]
+
+
+class NodeNetworkServiceReconAction(AbstractAction):
+    """Action which performs an NMAP network service recon (ping scan followed by port scan)."""
+
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
+        super().__init__(manager=manager)
+
+    def form_request(
+        self,
+        source_node: str,
+        target_ip_address: Union[str, List[str]],
+        target_protocol: Optional[Union[str, List[str]]] = None,
+        target_port: Optional[Union[str, List[str]]] = None,
+    ) -> List[str]:  # noqa
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
+        return [
+            "network",
+            "node",
+            source_node,
+            "application",
+            "NMAP",
+            "network_service_recon",
+            {"target_ip_address": target_ip_address, "target_port": target_port, "target_protocol": target_protocol},
+        ]
+
+
 class ActionManager:
     """Class which manages the action space for an agent."""
 
@@ -916,6 +984,9 @@ class ActionManager:
         "HOST_NIC_DISABLE": HostNICDisableAction,
         "NETWORK_PORT_ENABLE": NetworkPortEnableAction,
         "NETWORK_PORT_DISABLE": NetworkPortDisableAction,
+        "NODE_NMAP_PING_SCAN": NodeNMAPPingScanAction,
+        "NODE_NMAP_PORT_SCAN": NodeNMAPPortScanAction,
+        "NODE_NMAP_NETWORK_SERVICE_RECON": NodeNetworkServiceReconAction,
     }
     """Dictionary which maps action type strings to the corresponding action class."""
 

src/primaite/game/game.py

@@ -361,11 +361,6 @@ class PrimaiteGame:
                                 server_ip_address=IPv4Address(opt.get("server_ip")),
                                 server_password=opt.get("server_password"),
                                 payload=opt.get("payload", "ENCRYPT"),
-                                c2_beacon_p_of_success=float(opt.get("c2_beacon_p_of_success", "0.5")),
-                                target_scan_p_of_success=float(opt.get("target_scan_p_of_success", "0.1")),
-                                ransomware_encrypt_p_of_success=float(
-                                    opt.get("ransomware_encrypt_p_of_success", "0.1")
-                                ),
                             )
                     elif application_type == "DatabaseClient":
                         if "options" in application_cfg:

src/primaite/interface/request.py

@@ -3,7 +3,7 @@ from typing import Dict, ForwardRef, List, Literal, Union
 
 from pydantic import BaseModel, ConfigDict, StrictBool, validate_call
 
-RequestFormat = List[Union[str, int, float]]
+RequestFormat = List[Union[str, int, float, Dict]]
 
 RequestResponse = ForwardRef("RequestResponse")
 """This makes it possible to type-hint RequestResponse.from_bool return type."""

src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb

@@ -4,13 +4,15 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Customising red agents\n",
+    "# Customising UC2 Red Agents\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
     "\n",
     "This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
     "\n",
     "First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
     "\n",
-    "*(For a full explanation of the Data Manipulation scenario, check out the notebook `Data-Manipulation-E2E-Demonstration.ipynb`)*"
+    "*(For a full explanation of the Data Manipulation scenario, check out the data manipulation scenario notebook)*"
    ]
   },
   {

src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb

@@ -4,7 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Data Manipulation Scenario\n"
+    "# Data Manipulation Scenario\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
    ]
   },
   {
@@ -79,7 +81,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Reinforcement learning details"
+    "## Reinforcement learning details"
    ]
   },
   {
@@ -692,7 +694,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "venv",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },

src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb

@@ -4,7 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Getting information out of PrimAITE"
+    "# Getting information out of PrimAITE\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
    ]
   },
   {
@@ -160,7 +162,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.10.8"
   }
  },
  "nbformat": 4,

src/primaite/notebooks/Requests-and-Responses.ipynb

@@ -6,7 +6,9 @@
    "source": [
     "# Requests and Responses\n",
     "\n",
-    "Agents interact with the PrimAITE simulation via the Request system.\n"
+    "Agents interact with the PrimAITE simulation via the Request system.\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
    ]
   },
   {

src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb

@@ -4,7 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Train a Multi agent system using RLLIB\n",
+    "# Train a Multi agent system using RLLIB\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
     "\n",
     "This notebook will demonstrate how to use the `PrimaiteRayMARLEnv` to train a very basic system with two PPO agents."
    ]

src/primaite/notebooks/Training-an-RLLib-Agent.ipynb

@@ -4,7 +4,10 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Train a Single agent system using RLLib\n",
+    "# Train a Single agent system using RLLib\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
     "This notebook will demonstrate how to use PrimaiteRayEnv to train a basic PPO agent."
    ]
   },

src/primaite/notebooks/Training-an-SB3-Agent.ipynb

@@ -6,6 +6,8 @@
    "source": [
     "# Training an SB3 Agent\n",
     "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
     "This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file."
    ]
   },
@@ -166,7 +168,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "venv",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -180,7 +182,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.10.8"
   }
  },
  "nbformat": 4,

src/primaite/notebooks/Using-Episode-Schedules.ipynb

@@ -6,6 +6,8 @@
    "source": [
     "# Using Episode Schedules\n",
     "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
     "PrimAITE supports the ability to use different variations on a scenario at different episodes. This can be used to increase \n",
     "domain randomisation to prevent overfitting, or to set up curriculum learning to train agents to perform more complicated tasks.\n",
     "\n",

src/primaite/notebooks/multi-processing.ipynb

@@ -4,7 +4,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Simple multi-processing demo using SubprocVecEnv from SB3"
+    "# Simple multi-processing demonstration\n",
+    "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
+    "This notebook uses SubprocVecEnv from SB3."
    ]
   },
   {
@@ -139,7 +143,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

src/primaite/session/environment.py

@@ -38,6 +38,8 @@ class PrimaiteGymEnv(gymnasium.Env):
         """Name of the RL agent. Since there should only be one RL agent we can just pull the first and only key."""
         self.episode_counter: int = 0
         """Current episode number."""
+        self.total_reward_per_episode: Dict[int, float] = {}
+        """Average rewards of agents per episode."""
 
     @property
     def agent(self) -> ProxyAgent:
@@ -90,6 +92,8 @@ class PrimaiteGymEnv(gymnasium.Env):
             f"Resetting environment, episode {self.episode_counter}, "
             f"avg. reward: {self.agent.reward_function.total_reward}"
         )
+        self.total_reward_per_episode[self.episode_counter] = self.agent.reward_function.total_reward
+
         if self.io.settings.save_agent_actions:
             all_agent_actions = {name: agent.history for name, agent in self.game.agents.items()}
             self.io.write_agent_log(agent_actions=all_agent_actions, episode=self.episode_counter)

src/primaite/simulator/_package_data/create-simulation_demo.ipynb

@@ -6,7 +6,9 @@
    "source": [
     "# Build a simulation using the Python API\n",
     "\n",
-    "Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n"
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
+    "Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects."
    ]
   },
   {

src/primaite/simulator/_package_data/network_simulator_demo.ipynb

@@ -7,6 +7,8 @@
    "source": [
     "# PrimAITE Router Simulation Demo\n",
     "\n",
+    "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
+    "\n",
     "This demo uses a modified version of the ARCD Use Case 2 Network (seen below) to demonstrate the capabilities of the Network simulator in PrimAITE."
    ]
   },

src/primaite/simulator/core.py

@@ -222,7 +222,7 @@ class SimComponent(BaseModel):
         return state
 
     @validate_call
-    def apply_request(self, request: RequestFormat, context: Dict = {}) -> RequestResponse:
+    def apply_request(self, request: RequestFormat, context: Optional[Dict] = None) -> RequestResponse:
         """
         Apply a request to a simulation component. Request data is passed in as a 'namespaced' list of strings.
 
@@ -240,6 +240,8 @@ class SimComponent(BaseModel):
         :param: context: Dict containing context for requests
         :type context: Dict
         """
+        if not context:
+            context = None
         if self._request_manager is None:
             return
         return self._request_manager(request, context)

src/primaite/simulator/network/hardware/base.py

@@ -29,6 +29,7 @@ from primaite.simulator.network.nmne import (
     NMNE_CAPTURE_KEYWORDS,
 )
 from primaite.simulator.network.transmission.data_link_layer import Frame
+from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.core.packet_capture import PacketCapture
 from primaite.simulator.system.core.session_manager import SessionManager
@@ -37,6 +38,7 @@ from primaite.simulator.system.core.sys_log import SysLog
 from primaite.simulator.system.processes.process import Process
 from primaite.simulator.system.services.service import Service
 from primaite.simulator.system.software import IOSoftware
+from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
 from primaite.utils.validators import IPV4Address
 
 IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
@@ -108,10 +110,14 @@ class NetworkInterface(SimComponent, ABC):
     nmne: Dict = Field(default_factory=lambda: {})
     "A dict containing details of the number of malicious network events captured."
 
+    traffic: Dict = Field(default_factory=lambda: {})
+    "A dict containing details of the inbound and outbound traffic by port and protocol."
+
     def setup_for_episode(self, episode: int):
         """Reset the original state of the SimComponent."""
         super().setup_for_episode(episode=episode)
         self.nmne = {}
+        self.traffic = {}
         if episode and self.pcap and SIM_OUTPUT.save_pcap_logs:
             self.pcap.current_episode = episode
             self.pcap.setup_logger()
@@ -147,6 +153,7 @@ class NetworkInterface(SimComponent, ABC):
         )
         if CAPTURE_NMNE:
             state.update({"nmne": {k: v for k, v in self.nmne.items()}})
+        state.update({"traffic": convert_dict_enum_keys_to_enum_values(self.traffic)})
         return state
 
     @abstractmethod
@@ -237,6 +244,47 @@ class NetworkInterface(SimComponent, ABC):
                 # Increment a generic counter if keyword capturing is not enabled
                 keyword_level["*"] = keyword_level.get("*", 0) + 1
 
+    def _capture_traffic(self, frame: Frame, inbound: bool = True):
+        """
+        Capture traffic statistics at the Network Interface.
+
+        :param frame: The network frame containing the traffic data.
+        :type frame: Frame
+        :param inbound: Flag indicating if the traffic is inbound or outbound. Defaults to True.
+        :type inbound: bool
+        """
+        # Determine the direction of the traffic
+        direction = "inbound" if inbound else "outbound"
+
+        # Initialize protocol and port variables
+        protocol = None
+        port = None
+
+        # Identify the protocol and port from the frame
+        if frame.tcp:
+            protocol = IPProtocol.TCP
+            port = frame.tcp.dst_port
+        elif frame.udp:
+            protocol = IPProtocol.UDP
+            port = frame.udp.dst_port
+        elif frame.icmp:
+            protocol = IPProtocol.ICMP
+
+        # Ensure the protocol is in the capture dict
+        if protocol not in self.traffic:
+            self.traffic[protocol] = {}
+
+        # Handle non-ICMP protocols that use ports
+        if protocol != IPProtocol.ICMP:
+            if port not in self.traffic[protocol]:
+                self.traffic[protocol][port] = {"inbound": 0, "outbound": 0}
+            self.traffic[protocol][port][direction] += frame.size
+        else:
+            # Handle ICMP protocol separately (ICMP does not use ports)
+            if not self.traffic[protocol]:
+                self.traffic[protocol] = {"inbound": 0, "outbound": 0}
+            self.traffic[protocol][direction] += frame.size
+
     @abstractmethod
     def send_frame(self, frame: Frame) -> bool:
         """
@@ -246,6 +294,7 @@ class NetworkInterface(SimComponent, ABC):
         :return: A boolean indicating whether the frame was successfully sent.
         """
         self._capture_nmne(frame, inbound=False)
+        self._capture_traffic(frame, inbound=False)
 
     @abstractmethod
     def receive_frame(self, frame: Frame) -> bool:
@@ -256,6 +305,7 @@ class NetworkInterface(SimComponent, ABC):
         :return: A boolean indicating whether the frame was successfully received.
         """
         self._capture_nmne(frame, inbound=True)
+        self._capture_traffic(frame, inbound=True)
 
     def __str__(self) -> str:
         """
@@ -767,6 +817,24 @@ class Node(SimComponent):
         self.session_manager.software_manager = self.software_manager
         self._install_system_software()
 
+    def ip_is_network_interface(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool:
+        """
+        Checks if a given IP address belongs to any of the nodes interfaces.
+
+        :param ip_address: The IP address to check.
+        :param enabled_only: If True, only considers enabled network interfaces.
+        :return: True if the IP address is assigned to one of the nodes interfaces; False otherwise.
+        """
+        for network_interface in self.network_interface.values():
+            if not hasattr(network_interface, "ip_address"):
+                continue
+            if network_interface.ip_address == ip_address:
+                if enabled_only:
+                    return network_interface.enabled
+                else:
+                    return True
+        return False
+
     def setup_for_episode(self, episode: int):
         """Reset the original state of the SimComponent."""
         super().setup_for_episode(episode=episode)

src/primaite/simulator/network/hardware/nodes/host/host_node.py

@@ -8,6 +8,8 @@ from primaite import getLogger
 from primaite.simulator.network.hardware.base import IPWiredNetworkInterface, Link, Node
 from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
 from primaite.simulator.network.transmission.data_link_layer import Frame
+from primaite.simulator.system.applications.application import ApplicationOperatingState
+from primaite.simulator.system.applications.nmap import NMAP
 from primaite.simulator.system.applications.web_browser import WebBrowser
 from primaite.simulator.system.services.arp.arp import ARP, ARPPacket
 from primaite.simulator.system.services.dns.dns_client import DNSClient
@@ -303,6 +305,7 @@ class HostNode(Node):
         "DNSClient": DNSClient,
         "NTPClient": NTPClient,
         "WebBrowser": WebBrowser,
+        "NMAP": NMAP,
     }
     """List of system software that is automatically installed on nodes."""
 
@@ -315,6 +318,16 @@ class HostNode(Node):
         super().__init__(**kwargs)
         self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
 
+    @property
+    def nmap(self) -> Optional[NMAP]:
+        """
+        Return the NMAP application installed on the Node.
+
+        :return: NMAP application installed on the Node.
+        :rtype: Optional[NMAP]
+        """
+        return self.software_manager.software.get("NMAP")
+
     @property
     def arp(self) -> Optional[ARP]:
         """
@@ -366,8 +379,15 @@ class HostNode(Node):
         elif frame.udp:
             dst_port = frame.udp.dst_port
 
+        can_accept_nmap = False
+        if self.software_manager.software.get("NMAP"):
+            if self.software_manager.software["NMAP"].operating_state == ApplicationOperatingState.RUNNING:
+                can_accept_nmap = True
+
+        accept_nmap = can_accept_nmap and frame.payload.__class__.__name__ == "PortScanPayload"
+
         accept_frame = False
-        if frame.icmp or dst_port in self.software_manager.get_open_ports():
+        if frame.icmp or dst_port in self.software_manager.get_open_ports() or accept_nmap:
             # accept the frame as the port is open or if it's an ICMP frame
             accept_frame = True
 

src/primaite/simulator/network/hardware/nodes/network/router.py

@@ -19,6 +19,7 @@ from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
 from primaite.simulator.network.transmission.data_link_layer import Frame
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.applications.nmap import NMAP
 from primaite.simulator.system.core.session_manager import SessionManager
 from primaite.simulator.system.core.sys_log import SysLog
 from primaite.simulator.system.services.arp.arp import ARP
@@ -1239,6 +1240,7 @@ class Router(NetworkNode):
         icmp.router = self
         self.software_manager.install(RouterARP)
         self.arp.router = self
+        self.software_manager.install(NMAP)
 
     def _set_default_acl(self):
         """

src/primaite/simulator/system/applications/database_client.py

@@ -271,9 +271,16 @@ class DatabaseClient(Application):
 
         Calls disconnect on all client connections to ensure that both client and server connections are killed.
         """
-        while self.client_connections.values():
-            client_connection = self.client_connections[next(iter(self.client_connections.keys()))]
-            client_connection.disconnect()
+        while self.client_connections:
+            conn_key = next(iter(self.client_connections.keys()))
+            conn_obj: DatabaseClientConnection = self.client_connections[conn_key]
+            conn_obj.disconnect()
+            if conn_obj.is_active or conn_key in self.client_connections:
+                self.sys_log.error(
+                    "Attempted to uninstall database client but could not drop active connections. "
+                    "Forcing uninstall anyway."
+                )
+                self.client_connections.pop(conn_key, None)
         super().uninstall()
 
     def get_new_connection(self) -> Optional[DatabaseClientConnection]:

src/primaite/simulator/system/applications/nmap.py

@@ -0,0 +1,452 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from ipaddress import IPv4Address, IPv4Network
+from typing import Any, Dict, Final, List, Optional, Set, Tuple, Union
+
+from prettytable import PrettyTable
+from pydantic import validate_call
+
+from primaite.interface.request import RequestResponse
+from primaite.simulator.core import RequestManager, RequestType, SimComponent
+from primaite.simulator.network.transmission.network_layer import IPProtocol
+from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.applications.application import Application
+from primaite.utils.validators import IPV4Address
+
+
+class PortScanPayload(SimComponent):
+    """
+    A class representing the payload for a port scan.
+
+    :ivar ip_address: The target IP address for the port scan.
+    :ivar port: The target port for the port scan.
+    :ivar protocol: The protocol used for the port scan.
+    :ivar request:Flag to indicate whether this is a request or not.
+    """
+
+    ip_address: IPV4Address
+    port: Port
+    protocol: IPProtocol
+    request: bool = True
+
+    def describe_state(self) -> Dict:
+        """
+        Describe the state of the port scan payload.
+
+        :return: A dictionary representation of the port scan payload state.
+        :rtype: Dict
+        """
+        state = super().describe_state()
+        state["ip_address"] = str(self.ip_address)
+        state["port"] = self.port.value
+        state["protocol"] = self.protocol.value
+        state["request"] = self.request
+
+        return state
+
+
+class NMAP(Application):
+    """
+    A class representing the NMAP application for network scanning.
+
+    NMAP is a network scanning tool used to discover hosts and services on a network. It provides functionalities such
+    as ping scans to discover active hosts and port scans to detect open ports on those hosts.
+    """
+
+    _active_port_scans: Dict[str, PortScanPayload] = {}
+    _port_scan_responses: Dict[str, PortScanPayload] = {}
+
+    _PORT_SCAN_TYPE_MAP: Final[Dict[Tuple[bool, bool], str]] = {
+        (True, True): "Box",
+        (True, False): "Horizontal",
+        (False, True): "Vertical",
+        (False, False): "Port",
+    }
+
+    def __init__(self, **kwargs):
+        kwargs["name"] = "NMAP"
+        kwargs["port"] = Port.NONE
+        kwargs["protocol"] = IPProtocol.NONE
+        super().__init__(**kwargs)
+
+    def _can_perform_network_action(self) -> bool:
+        """
+        Checks if the NMAP application can perform outbound network actions.
+
+        This is done by checking the parent application can_per_action functionality. Then checking if there is an
+        enabled NIC that can be used for outbound traffic.
+
+        :return: True if outbound network actions can be performed, otherwise False.
+        """
+        if not super()._can_perform_action():
+            return False
+
+        for nic in self.software_manager.node.network_interface.values():
+            if nic.enabled:
+                return True
+        return False
+
+    def _init_request_manager(self) -> RequestManager:
+        def _ping_scan_action(request: List[Any], context: Any) -> RequestResponse:
+            results = self.ping_scan(target_ip_address=request[0]["target_ip_address"], json_serializable=True)
+            if not self._can_perform_network_action():
+                return RequestResponse.from_bool(False)
+            return RequestResponse(
+                status="success",
+                data={"live_hosts": results},
+            )
+
+        def _port_scan_action(request: List[Any], context: Any) -> RequestResponse:
+            results = self.port_scan(**request[0], json_serializable=True)
+            if not self._can_perform_network_action():
+                return RequestResponse.from_bool(False)
+            return RequestResponse(
+                status="success",
+                data=results,
+            )
+
+        def _network_service_recon_action(request: List[Any], context: Any) -> RequestResponse:
+            results = self.network_service_recon(**request[0], json_serializable=True)
+            if not self._can_perform_network_action():
+                return RequestResponse.from_bool(False)
+            return RequestResponse(
+                status="success",
+                data=results,
+            )
+
+        rm = RequestManager()
+
+        rm.add_request(
+            name="ping_scan",
+            request_type=RequestType(func=_ping_scan_action),
+        )
+
+        rm.add_request(
+            name="port_scan",
+            request_type=RequestType(func=_port_scan_action),
+        )
+
+        rm.add_request(
+            name="network_service_recon",
+            request_type=RequestType(func=_network_service_recon_action),
+        )
+
+        return rm
+
+    def describe_state(self) -> Dict:
+        """
+        Describe the state of the NMAP application.
+
+        :return: A dictionary representation of the NMAP application's state.
+        :rtype: Dict
+        """
+        return super().describe_state()
+
+    @staticmethod
+    def _explode_ip_address_network_array(
+        target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
+    ) -> Set[IPv4Address]:
+        """
+        Explode a mixed array of IP addresses and networks into a set of individual IP addresses.
+
+        This method takes a combination of single and lists of IPv4 addresses and IPv4 networks, expands any networks
+        into their constituent subnet useable IP addresses, and returns a set of unique IP addresses. Broadcast and
+        network addresses are excluded from the result.
+
+        :param target_ip_address: A single or list of IPv4 addresses and networks.
+        :type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
+        :return: A set of unique IPv4 addresses expanded from the input.
+        :rtype: Set[IPv4Address]
+        """
+        if isinstance(target_ip_address, IPv4Address) or isinstance(target_ip_address, IPv4Network):
+            target_ip_address = [target_ip_address]
+        ip_addresses: List[IPV4Address] = []
+        for ip_address in target_ip_address:
+            if isinstance(ip_address, IPv4Network):
+                ip_addresses += [
+                    ip
+                    for ip in ip_address.hosts()
+                    if not ip == ip_address.broadcast_address and not ip == ip_address.network_address
+                ]
+            else:
+                ip_addresses.append(ip_address)
+        return set(ip_addresses)
+
+    @validate_call()
+    def ping_scan(
+        self,
+        target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
+        show: bool = True,
+        show_online_only: bool = True,
+        json_serializable: bool = False,
+    ) -> Union[List[IPV4Address], List[str]]:
+        """
+        Perform a ping scan on the target IP address(es).
+
+        :param target_ip_address: The target IP address(es) or network(s) for the ping scan.
+        :type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
+        :param show: Flag indicating whether to display the scan results. Defaults to True.
+        :type show: bool
+        :param show_online_only: Flag indicating whether to show only the online hosts. Defaults to True.
+        :type show_online_only: bool
+        :param json_serializable: Flag indicating whether the return value should be json serializable. Defaults to
+            False.
+        :type json_serializable: bool
+
+        :return: A list of active IP addresses that responded to the ping.
+        :rtype: Union[List[IPV4Address], List[str]]
+        """
+        active_nodes = []
+        if show:
+            table = PrettyTable(["IP Address", "Can Ping"])
+            table.align = "l"
+            table.title = f"{self.software_manager.node.hostname} NMAP Ping Scan"
+
+        ip_addresses = self._explode_ip_address_network_array(target_ip_address)
+
+        for ip_address in ip_addresses:
+            # Prevent ping scan on this node
+            if self.software_manager.node.ip_is_network_interface(ip_address=ip_address):
+                continue
+            can_ping = self.software_manager.icmp.ping(ip_address)
+            if can_ping:
+                active_nodes.append(ip_address if not json_serializable else str(ip_address))
+            if show and (can_ping or not show_online_only):
+                table.add_row([ip_address, can_ping])
+        if show:
+            print(table.get_string(sortby="IP Address"))
+        return active_nodes
+
+    def _determine_port_scan_type(self, target_ip_addresses: List[IPV4Address], target_ports: List[Port]) -> str:
+        """
+        Determine the type of port scan based on the number of target IP addresses and ports.
+
+        :param target_ip_addresses: The list of target IP addresses.
+        :type target_ip_addresses: List[IPV4Address]
+        :param target_ports: The list of target ports.
+        :type target_ports: List[Port]
+
+        :return: The type of port scan.
+        :rtype: str
+        """
+        vertical_scan = len(target_ports) > 1
+        horizontal_scan = len(target_ip_addresses) > 1
+
+        return self._PORT_SCAN_TYPE_MAP[horizontal_scan, vertical_scan]
+
+    def _check_port_open_on_ip_address(
+        self,
+        ip_address: IPv4Address,
+        port: Port,
+        protocol: IPProtocol,
+        is_re_attempt: bool = False,
+        port_scan_uuid: Optional[str] = None,
+    ) -> bool:
+        """
+        Check if a port is open on a specific IP address.
+
+        :param ip_address: The target IP address.
+        :type ip_address: IPv4Address
+        :param port: The target port.
+        :type port: Port
+        :param protocol: The protocol used for the port scan.
+        :type protocol: IPProtocol
+        :param is_re_attempt: Flag indicating if this is a reattempt. Defaults to False.
+        :type is_re_attempt: bool
+        :param port_scan_uuid: The UUID of the port scan payload. Defaults to None.
+        :type port_scan_uuid: Optional[str]
+
+        :return: True if the port is open, False otherwise.
+        :rtype: bool
+        """
+        # The recursive base case
+        if is_re_attempt:
+            # Return True if a response has been received, otherwise return False
+            if port_scan_uuid in self._port_scan_responses:
+                self._port_scan_responses.pop(port_scan_uuid)
+                return True
+            return False
+
+        # Send the port scan request
+        payload = PortScanPayload(ip_address=ip_address, port=port, protocol=protocol)
+        self._active_port_scans[payload.uuid] = payload
+        self.sys_log.info(
+            f"{self.name}: Sending port scan request over {payload.protocol.name} on port {payload.port.value} "
+            f"({payload.port.name}) to {payload.ip_address}"
+        )
+        self.software_manager.send_payload_to_session_manager(
+            payload=payload, dest_ip_address=ip_address, src_port=port, dest_port=port, ip_protocol=protocol
+        )
+
+        # Recursively call this function with as a reattempt
+        return self._check_port_open_on_ip_address(
+            ip_address=ip_address, port=port, protocol=protocol, is_re_attempt=True, port_scan_uuid=payload.uuid
+        )
+
+    def _process_port_scan_response(self, payload: PortScanPayload):
+        """
+        Process the response to a port scan request.
+
+        :param payload: The port scan payload received in response.
+        :type payload: PortScanPayload
+        """
+        if payload.uuid in self._active_port_scans:
+            self._active_port_scans.pop(payload.uuid)
+            self._port_scan_responses[payload.uuid] = payload
+            self.sys_log.info(
+                f"{self.name}: Received port scan response from {payload.ip_address} on port {payload.port.value} "
+                f"({payload.port.name}) over {payload.protocol.name}"
+            )
+
+    def _process_port_scan_request(self, payload: PortScanPayload, session_id: str) -> None:
+        """
+        Process a port scan request.
+
+        :param payload: The port scan payload received in the request.
+        :type payload: PortScanPayload
+        :param session_id: The session ID for the port scan request.
+        :type session_id: str
+        """
+        if self.software_manager.check_port_is_open(port=payload.port, protocol=payload.protocol):
+            payload.request = False
+            self.sys_log.info(
+                f"{self.name}: Responding to port scan request for port {payload.port.value} "
+                f"({payload.port.name}) over {payload.protocol.name}",
+                True,
+            )
+            self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
+
+    @validate_call()
+    def port_scan(
+        self,
+        target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
+        target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
+        target_port: Optional[Union[Port, List[Port]]] = None,
+        show: bool = True,
+        json_serializable: bool = False,
+    ) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
+        """
+        Perform a port scan on the target IP address(es).
+
+        :param target_ip_address: The target IP address(es) or network(s) for the port scan.
+        :type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
+        :param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
+        :type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
+        :param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
+        :type target_port: Optional[Union[Port, List[Port]]]
+        :param show: Flag indicating whether to display the scan results. Defaults to True.
+        :type show: bool
+        :param json_serializable: Flag indicating whether the return value should be JSON serializable. Defaults to
+            False.
+        :type json_serializable: bool
+
+        :return: A dictionary mapping IP addresses to protocols and lists of open ports.
+        :rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
+        """
+        ip_addresses = self._explode_ip_address_network_array(target_ip_address)
+
+        if isinstance(target_port, Port):
+            target_port = [target_port]
+        elif target_port is None:
+            target_port = [port for port in Port if port not in {Port.NONE, Port.UNUSED}]
+
+        if isinstance(target_protocol, IPProtocol):
+            target_protocol = [target_protocol]
+        elif target_protocol is None:
+            target_protocol = [IPProtocol.TCP, IPProtocol.UDP]
+
+        scan_type = self._determine_port_scan_type(list(ip_addresses), target_port)
+        active_ports = {}
+        if show:
+            table = PrettyTable(["IP Address", "Port", "Name", "Protocol"])
+            table.align = "l"
+            table.title = f"{self.software_manager.node.hostname} NMAP Port Scan ({scan_type})"
+        self.sys_log.info(f"{self.name}: Starting port scan")
+        for ip_address in ip_addresses:
+            # Prevent port scan on this node
+            if self.software_manager.node.ip_is_network_interface(ip_address=ip_address):
+                continue
+            for protocol in target_protocol:
+                for port in set(target_port):
+                    port_open = self._check_port_open_on_ip_address(ip_address=ip_address, port=port, protocol=protocol)
+
+                    if port_open:
+                        table.add_row([ip_address, port.value, port.name, protocol.name])
+                        _ip_address = ip_address if not json_serializable else str(ip_address)
+                        _protocol = protocol if not json_serializable else protocol.value
+                        _port = port if not json_serializable else port.value
+                        if _ip_address not in active_ports:
+                            active_ports[_ip_address] = dict()
+                        if _protocol not in active_ports[_ip_address]:
+                            active_ports[_ip_address][_protocol] = []
+                        active_ports[_ip_address][_protocol].append(_port)
+
+        if show:
+            print(table.get_string(sortby="IP Address"))
+
+        return active_ports
+
+    def network_service_recon(
+        self,
+        target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
+        target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
+        target_port: Optional[Union[Port, List[Port]]] = None,
+        show: bool = True,
+        show_online_only: bool = True,
+        json_serializable: bool = False,
+    ) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
+        """
+        Perform a network service reconnaissance which includes a ping scan followed by a port scan.
+
+        This method combines the functionalities of a ping scan and a port scan to provide a comprehensive
+        overview of the services on the network. It first identifies active hosts in the target IP range by performing
+        a ping scan. Once the active hosts are identified, it performs a port scan on these hosts to identify open
+        ports and running services. This two-step process ensures that the port scan is performed only on live hosts,
+        optimising the scanning process and providing accurate results.
+
+        :param target_ip_address: The target IP address(es) or network(s) for the port scan.
+        :type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
+        :param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
+        :type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
+        :param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
+        :type target_port: Optional[Union[Port, List[Port]]]
+        :param show: Flag indicating whether to display the scan results. Defaults to True.
+        :type show: bool
+        :param show_online_only: Flag indicating whether to show only the online hosts. Defaults to True.
+        :type show_online_only: bool
+        :param json_serializable: Flag indicating whether the return value should be JSON serializable. Defaults to
+            False.
+        :type json_serializable: bool
+
+        :return: A dictionary mapping IP addresses to protocols and lists of open ports.
+        :rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
+        """
+        ping_scan_results = self.ping_scan(
+            target_ip_address=target_ip_address, show=show, show_online_only=show_online_only, json_serializable=False
+        )
+        return self.port_scan(
+            target_ip_address=ping_scan_results,
+            target_protocol=target_protocol,
+            target_port=target_port,
+            show=show,
+            json_serializable=json_serializable,
+        )
+
+    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
+        """
+        Receive and process a payload.
+
+        :param payload: The payload to be processed.
+        :type payload: Any
+        :param session_id: The session ID associated with the payload.
+        :type session_id: str
+
+        :return: True if the payload was successfully processed, False otherwise.
+        :rtype: bool
+        """
+        if isinstance(payload, PortScanPayload):
+            if payload.request:
+                self._process_port_scan_request(payload=payload, session_id=session_id)
+            else:
+                self._process_port_scan_response(payload=payload)
+
+        return True

src/primaite/simulator/system/applications/red_applications/ransomware_script.py

@@ -1,9 +1,7 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
-from enum import IntEnum
 from ipaddress import IPv4Address
 from typing import Dict, Optional
 
-from primaite.game.science import simulate_trial
 from primaite.interface.request import RequestResponse
 from primaite.simulator.core import RequestManager, RequestType
 from primaite.simulator.network.transmission.network_layer import IPProtocol
@@ -12,43 +10,10 @@ from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
 
 
-class RansomwareAttackStage(IntEnum):
-    """
-    Enumeration representing different attack stages of the ransomware script.
-
-    This enumeration defines the various stages a data manipulation attack can be in during its lifecycle
-    in the simulation.
-    Each stage represents a specific phase in the attack process.
-    """
-
-    NOT_STARTED = 0
-    "Indicates that the attack has not started yet."
-    DOWNLOAD = 1
-    "Installing the Encryption Script - Testing"
-    INSTALL = 2
-    "The stage where logon procedures are simulated."
-    ACTIVATE = 3
-    "Operating Status Changes"
-    PROPAGATE = 4
-    "Represents the stage of performing a horizontal port scan on the target."
-    COMMAND_AND_CONTROL = 5
-    "Represents the stage of setting up a rely C2 Beacon (Not Implemented)"
-    PAYLOAD = 6
-    "Stage of actively attacking the target."
-    SUCCEEDED = 7
-    "Indicates the attack has been successfully completed."
-    FAILED = 8
-    "Signifies that the attack has failed."
-
-
 class RansomwareScript(Application):
     """Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
 
-    :ivar payload: The attack stage query payload. (Default Corrupt)
-    :ivar target_scan_p_of_success: The probability of success for the target scan stage.
-    :ivar c2_beacon_p_of_success: The probability of success for the c2_beacon stage
-    :ivar ransomware_encrypt_p_of_success: The probability of success for the ransomware 'attack' (encrypt) stage.
-    :ivar repeat: Whether to repeat attacking once finished.
+    :ivar payload: The attack stage query payload. (Default ENCRYPT)
     """
 
     server_ip_address: Optional[IPv4Address] = None
@@ -57,16 +22,6 @@ class RansomwareScript(Application):
     """Password required to access the database."""
     payload: Optional[str] = "ENCRYPT"
     "Payload String for the payload stage"
-    target_scan_p_of_success: float = 0.9
-    "Probability of the target scan succeeding: Default 0.9"
-    c2_beacon_p_of_success: float = 0.9
-    "Probability of the c2 beacon setup stage succeeding: Default 0.9"
-    ransomware_encrypt_p_of_success: float = 0.9
-    "Probability of the ransomware attack succeeding: Default 0.9"
-    repeat: bool = False
-    "If true, the Denial of Service bot will keep performing the attack."
-    attack_stage: RansomwareAttackStage = RansomwareAttackStage.NOT_STARTED
-    "The ransomware attack stage. See RansomwareAttackStage Class"
 
     def __init__(self, **kwargs):
         kwargs["name"] = "RansomwareScript"
@@ -91,7 +46,7 @@ class RansomwareScript(Application):
     @property
     def _host_db_client(self) -> DatabaseClient:
         """Return the database client that is installed on the same machine as the Ransomware Script."""
-        db_client = self.software_manager.software.get("DatabaseClient")
+        db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
         if db_client is None:
             self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
         return db_client
@@ -109,16 +64,6 @@ class RansomwareScript(Application):
         )
         return rm
 
-    def _activate(self):
-        """
-        Simulate the install process as the initial stage of the attack.
-
-        Advances the attack stage to 'ACTIVATE' attack state.
-        """
-        if self.attack_stage == RansomwareAttackStage.INSTALL:
-            self.sys_log.info(f"{self.name}: Activated!")
-            self.attack_stage = RansomwareAttackStage.ACTIVATE
-
     def run(self) -> bool:
         """Calls the parent classes execute method before starting the application loop."""
         super().run()
@@ -134,20 +79,9 @@ class RansomwareScript(Application):
             return False
         if self.server_ip_address and self.payload:
             self.sys_log.info(f"{self.name}: Running")
-            self.attack_stage = RansomwareAttackStage.NOT_STARTED
-            self._local_download()
-            self._install()
-            self._activate()
-            self._perform_target_scan()
-            self._setup_beacon()
-            self._perform_ransomware_encrypt()
-
-            if self.repeat and self.attack_stage in (
-                RansomwareAttackStage.SUCCEEDED,
-                RansomwareAttackStage.FAILED,
-            ):
-                self.attack_stage = RansomwareAttackStage.NOT_STARTED
-            return True
+            if self._perform_ransomware_encrypt():
+                return True
+            return False
         else:
             self.sys_log.warning(f"{self.name}: Failed to start as it requires both a target_ip_address and payload.")
             return False
@@ -157,10 +91,6 @@ class RansomwareScript(Application):
         server_ip_address: IPv4Address,
         server_password: Optional[str] = None,
         payload: Optional[str] = None,
-        target_scan_p_of_success: Optional[float] = None,
-        c2_beacon_p_of_success: Optional[float] = None,
-        ransomware_encrypt_p_of_success: Optional[float] = None,
-        repeat: bool = True,
     ):
         """
         Configure the Ransomware Script to communicate with a DatabaseService.
@@ -168,10 +98,6 @@ class RansomwareScript(Application):
         :param server_ip_address: The IP address of the Node the DatabaseService is on.
         :param server_password: The password on the DatabaseService.
         :param payload: The attack stage query (Encrypt / Delete)
-        :param target_scan_p_of_success: The probability of success for the target scan stage.
-        :param c2_beacon_p_of_success: The probability of success for the c2_beacon stage
-        :param ransomware_encrypt_p_of_success: The probability of success for the ransomware 'attack' (encrypt) stage.
-        :param repeat: Whether to repeat attacking once finished.
         """
         if server_ip_address:
             self.server_ip_address = server_ip_address
@@ -179,74 +105,15 @@ class RansomwareScript(Application):
             self.server_password = server_password
         if payload:
             self.payload = payload
-        if target_scan_p_of_success:
-            self.target_scan_p_of_success = target_scan_p_of_success
-        if c2_beacon_p_of_success:
-            self.c2_beacon_p_of_success = c2_beacon_p_of_success
-        if ransomware_encrypt_p_of_success:
-            self.ransomware_encrypt_p_of_success = ransomware_encrypt_p_of_success
-        if repeat:
-            self.repeat = repeat
         self.sys_log.info(
-            f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
-            f"{repeat=}."
+            f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}."
         )
 
-    def _install(self):
-        """
-        Simulate the install stage in the kill-chain.
-
-        Advances the attack stage to 'ACTIVATE' if successful.
-
-        From this attack stage onwards.
-        the ransomware application is now visible from this point onwardin the observation space.
-        """
-        if self.attack_stage == RansomwareAttackStage.DOWNLOAD:
-            self.sys_log.info(f"{self.name}: Malware installed on the local file system")
-            downloads_folder = self.file_system.get_folder(folder_name="downloads")
-            ransomware_file = downloads_folder.get_file(file_name="ransom_script.pdf")
-            ransomware_file.num_access += 1
-            self.attack_stage = RansomwareAttackStage.INSTALL
-
-    def _setup_beacon(self):
-        """
-        Simulates setting up a c2 beacon; currently a pseudo step for increasing red variance.
-
-        Advances the attack stage to 'COMMAND AND CONTROL` if successful.
-
-        :param p_of_sucess: Probability of a successful c2 setup (Advancing this step),
-                            by default the success rate is 0.5
-        """
-        if self.attack_stage == RansomwareAttackStage.PROPAGATE:
-            self.sys_log.info(f"{self.name} Attempting to set up C&C Beacon - Scan 1/2")
-            if simulate_trial(self.c2_beacon_p_of_success):
-                self.sys_log.info(f"{self.name} C&C Successful setup - Scan 2/2")
-                c2c_setup = True  # TODO Implement the c2c step via an FTP Application/Service
-                if c2c_setup:
-                    self.attack_stage = RansomwareAttackStage.COMMAND_AND_CONTROL
-
-    def _perform_target_scan(self):
-        """
-        Perform a simulated port scan to check for open SQL ports.
-
-        Advances the attack stage to `PROPAGATE` if successful.
-
-        :param p_of_success: Probability of successful port scan, by default 0.1.
-        """
-        if self.attack_stage == RansomwareAttackStage.ACTIVATE:
-            # perform a port scan to identify that the SQL port is open on the server
-            self.sys_log.info(f"{self.name}: Scanning for vulnerable databases - Scan 0/2")
-            if simulate_trial(self.target_scan_p_of_success):
-                self.sys_log.info(f"{self.name}: Found a target database! Scan 1/2")
-                port_is_open = True  # TODO Implement a NNME Triggering scan as a seperate Red Application
-                if port_is_open:
-                    self.attack_stage = RansomwareAttackStage.PROPAGATE
-
     def attack(self) -> bool:
         """Perform the attack steps after opening the application."""
+        self.run()
         if not self._can_perform_action():
             self.sys_log.warning("Ransomware application is unable to perform it's actions.")
-            self.run()
         self.num_executions += 1
         return self._application_loop()
 
@@ -255,57 +122,30 @@ class RansomwareScript(Application):
         self._db_connection = self._host_db_client.get_new_connection()
         return True if self._db_connection else False
 
-    def _perform_ransomware_encrypt(self):
+    def _perform_ransomware_encrypt(self) -> bool:
         """
         Execute the Ransomware Encrypt payload on the target.
 
         Advances the attack stage to `COMPLETE` if successful, or 'FAILED' if unsuccessful.
-        :param p_of_success: Probability of successfully performing ransomware encryption, by default 0.1.
         """
         if self._host_db_client is None:
             self.sys_log.info(f"{self.name}: Failed to connect to db_client - Ransomware Script")
-            self.attack_stage = RansomwareAttackStage.FAILED
-            return
+            return False
 
         self._host_db_client.server_ip_address = self.server_ip_address
         self._host_db_client.server_password = self.server_password
-        if self.attack_stage == RansomwareAttackStage.COMMAND_AND_CONTROL:
-            if simulate_trial(self.ransomware_encrypt_p_of_success):
-                self.sys_log.info(f"{self.name}: Attempting to launch payload")
-                if not self._db_connection:
-                    self._establish_db_connection()
-                if self._db_connection:
-                    attack_successful = self._db_connection.query(self.payload)
-                    self.sys_log.info(f"{self.name} Payload delivered: {self.payload}")
-                    if attack_successful:
-                        self.sys_log.info(f"{self.name}: Payload Successful")
-                        self.attack_stage = RansomwareAttackStage.SUCCEEDED
-                    else:
-                        self.sys_log.info(f"{self.name}: Payload failed")
-                        self.attack_stage = RansomwareAttackStage.FAILED
+        self.sys_log.info(f"{self.name}: Attempting to launch payload")
+        if not self._db_connection:
+            self._establish_db_connection()
+        if self._db_connection:
+            attack_successful = self._db_connection.query(self.payload)
+            self.sys_log.info(f"{self.name} Payload delivered: {self.payload}")
+            if attack_successful:
+                self.sys_log.info(f"{self.name}: Payload Successful")
+                return True
+            else:
+                self.sys_log.info(f"{self.name}: Payload failed")
+            return False
         else:
             self.sys_log.warning("Attack Attempted to launch too quickly")
-            self.attack_stage = RansomwareAttackStage.FAILED
-
-    def _local_download(self):
-        """Downloads itself via the onto the local file_system."""
-        if self.attack_stage == RansomwareAttackStage.NOT_STARTED:
-            if self._local_download_verify():
-                self.attack_stage = RansomwareAttackStage.DOWNLOAD
-            else:
-                self.sys_log.info("Malware failed to create a installation location")
-                self.attack_stage = RansomwareAttackStage.FAILED
-        else:
-            self.sys_log.info("Malware failed to download")
-            self.attack_stage = RansomwareAttackStage.FAILED
-
-    def _local_download_verify(self) -> bool:
-        """Verifies a download location - Creates one if needed."""
-        for folder in self.file_system.folders:
-            if self.file_system.folders[folder].name == "downloads":
-                self.file_system.num_file_creations += 1
-                return True
-
-        self.file_system.create_folder("downloads")
-        self.file_system.create_file(folder_name="downloads", file_name="ransom_script.pdf")
-        return True
+            return False

src/primaite/simulator/system/core/software_manager.py

@@ -79,6 +79,31 @@ class SoftwareManager:
                 open_ports.append(software.port)
         return open_ports
 
+    def check_port_is_open(self, port: Port, protocol: IPProtocol) -> bool:
+        """
+        Check if a specific port is open and running a service using the specified protocol.
+
+        This method iterates through all installed software on the node and checks if any of them
+        are using the specified port and protocol and are currently in a running state. It returns True if any software
+        is found running on the specified port and protocol, otherwise False.
+
+
+        :param port: The port to check.
+        :type port: Port
+        :param protocol: The protocol to check (e.g., TCP, UDP).
+        :type protocol: IPProtocol
+        :return: True if the port is open and a service is running on it using the specified protocol, False otherwise.
+        :rtype: bool
+        """
+        for software in self.software.values():
+            if (
+                software.port == port
+                and software.protocol == protocol
+                and software.operating_state in {ApplicationOperatingState.RUNNING, ServiceOperatingState.RUNNING}
+            ):
+                return True
+        return False
+
     def install(self, software_class: Type[IOSoftwareClass]):
         """
         Install an Application or Service.
@@ -151,6 +176,7 @@ class SoftwareManager:
         self,
         payload: Any,
         dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
+        src_port: Optional[Port] = None,
         dest_port: Optional[Port] = None,
         ip_protocol: IPProtocol = IPProtocol.TCP,
         session_id: Optional[str] = None,
@@ -171,6 +197,7 @@ class SoftwareManager:
         return self.session_manager.receive_payload_from_software_manager(
             payload=payload,
             dst_ip_address=dest_ip_address,
+            src_port=src_port,
             dst_port=dest_port,
             ip_protocol=ip_protocol,
             session_id=session_id,
@@ -191,6 +218,9 @@ class SoftwareManager:
         :param payload: The payload being received.
         :param session: The transport session the payload originates from.
         """
+        if payload.__class__.__name__ == "PortScanPayload":
+            self.software.get("NMAP").receive(payload=payload, session_id=session_id)
+            return
         receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None)
         if receiver:
             receiver.receive(

src/primaite/utils/converters.py

@@ -0,0 +1,26 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from enum import Enum
+from typing import Any, Dict
+
+
+def convert_dict_enum_keys_to_enum_values(d: Dict[Any, Any]) -> Dict[Any, Any]:
+    """
+    Convert dictionary keys from enums to their corresponding values.
+
+    :param d: dict
+        The dictionary with enum keys to be converted.
+    :return: dict
+        The dictionary with enum values as keys.
+    """
+    result = {}
+    for key, value in d.items():
+        if isinstance(key, Enum):
+            new_key = key.value
+        else:
+            new_key = key
+
+        if isinstance(value, dict):
+            result[new_key] = convert_dict_enum_keys_to_enum_values(value)
+        else:
+            result[new_key] = value
+    return result

src/primaite/utils/session_output_reader.py

@@ -3,7 +3,7 @@
 raise DeprecationWarning(
     "Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
 )
-# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 from pathlib import Path
 from typing import Any, Dict, Tuple, Union
 
@@ -12,16 +12,16 @@ from typing import Any, Dict, Tuple, Union
 import polars as pl
 
 
-def av_rewards_dict(av_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
+def total_rewards_dict(total_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
     """
     Read an average rewards per episode csv file and return as a dict.
 
     The dictionary keys are the episode number, and the values are the mean reward that episode.
 
-    :param av_rewards_csv_file: The average rewards per episode csv file path.
+    :param total_rewards_csv_file: The average rewards per episode csv file path.
     :return: The average rewards per episode csv as a dict.
     """
-    df_dict = pl.read_csv(av_rewards_csv_file).to_dict()
+    df_dict = pl.read_csv(total_rewards_csv_file).to_dict()
 
     return {int(v): df_dict["Average Reward"][i] for i, v in enumerate(df_dict["Episode"])}
 

src/primaite/utils/session_output_writer.py

@@ -3,7 +3,7 @@
 raise DeprecationWarning(
     "Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
 )
-# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 import csv
 from logging import Logger
 from typing import Final, List, Tuple, TYPE_CHECKING, Union
@@ -27,9 +27,9 @@ class SessionOutputWriter:
     Is used to write session outputs to csv file.
     """
 
-    _AV_REWARD_PER_EPISODE_HEADER: Final[List[str]] = [
+    _TOTAL_REWARD_PER_EPISODE_HEADER: Final[List[str]] = [
         "Episode",
-        "Average Reward",
+        "Total Reward",
     ]
 
     def __init__(
@@ -44,7 +44,7 @@ class SessionOutputWriter:
         :param env: PrimAITE gym environment.
         :type env: Primaite
         :param transaction_writer: If `true`, this will output a full account of every transaction taken by the agent.
-            If `false` it will output the average reward per episode, defaults to False
+            If `false` it will output the total reward per episode, defaults to False
         :type transaction_writer: bool, optional
         :param learning_session: Set to `true` to indicate that the current session is a training session. This
             determines the name of the folder which contains the final output csv. Defaults to True
@@ -57,7 +57,7 @@ class SessionOutputWriter:
         if self.transaction_writer:
             fn = f"all_transactions_{self._env.timestamp_str}.csv"
         else:
-            fn = f"average_reward_per_episode_{self._env.timestamp_str}.csv"
+            fn = f"total_reward_per_episode_{self._env.timestamp_str}.csv"
 
         self._csv_file_path: "Path"
         if self.learning_session:
@@ -95,7 +95,7 @@ class SessionOutputWriter:
         if isinstance(data, Transaction):
             header, data = data.as_csv_data()
         else:
-            header = self._AV_REWARD_PER_EPISODE_HEADER
+            header = self._TOTAL_REWARD_PER_EPISODE_HEADER
 
         if self._first_write:
             self._init_csv_writer()