Merge remote-tracking branch 'origin/dev' into feature/2769-implement-user-account-action-space

CHANGELOG.md

@@ -9,9 +9,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 -   Random Number Generator Seeding by specifying a random number seed in the config file.
 -   Implemented Terminal service class, providing a generic terminal simulation.
+-   Added `User`, `UserManager` and `UserSessionManager` to enable the creation of user accounts and login on Nodes.
+-   Added a `listen_on_ports` set in the `IOSoftware` class to enable software listening on ports in addition to the
+    main port they're assigned.
 
 ### Changed
--   Removed the install/uninstall methods in the node class and made the software manager install/uninstall handle all of their functionality.
+-   Updated `SoftwareManager` `install` and `uninstall` to handle all functionality that was being done at the `install`
+    and `uninstall` methods in the `Node` class.
+-   Updated the `receive_payload_from_session_manager` method in `SoftwareManager` so that it now sends a copy of the
+    payload to any software listening on the destination port of the `Frame`.
+
+### Removed
+-   Removed the `install` and `uninstall` methods in the `Node` class.
 
 
 ## [3.2.0] - 2024-07-18

benchmark/primaite_benchmark.py

@@ -5,7 +5,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Final, Tuple
 
-from report import build_benchmark_md_report
+from report import build_benchmark_md_report, md2pdf
 from stable_baselines3 import PPO
 
 import primaite
@@ -159,6 +159,13 @@ def run(
     learning_rate: float = 3e-4,
 ) -> None:
     """Run the PrimAITE benchmark."""
+    # generate report folder
+    v_str = f"v{primaite.__version__}"
+
+    version_result_dir = _RESULTS_ROOT / v_str
+    version_result_dir.mkdir(exist_ok=True, parents=True)
+    output_path = version_result_dir / f"PrimAITE {v_str} Benchmark Report.md"
+
     benchmark_start_time = datetime.now()
 
     session_metadata_dict = {}
@@ -193,6 +200,12 @@ def run(
         session_metadata=session_metadata_dict,
         config_path=data_manipulation_config_path(),
         results_root_path=_RESULTS_ROOT,
+        output_path=output_path,
+    )
+    md2pdf(
+        md_path=output_path,
+        pdf_path=str(output_path).replace(".md", ".pdf"),
+        css_path="static/styles.css",
     )
 
 

benchmark/report.py

@@ -2,6 +2,7 @@
 import json
 import sys
 from datetime import datetime
+from os import PathLike
 from pathlib import Path
 from typing import Dict, Optional
 
@@ -14,7 +15,7 @@ from utils import _get_system_info
 import primaite
 
 PLOT_CONFIG = {
-    "size": {"auto_size": False, "width": 1500, "height": 900},
+    "size": {"auto_size": False, "width": 800, "height": 640},
     "template": "plotly_white",
     "range_slider": False,
 }
@@ -144,6 +145,20 @@ def _plot_benchmark_metadata(
         yaxis={"title": "Total Reward"},
         title=title,
     )
+    fig.update_layout(
+        legend=dict(
+            yanchor="top",
+            y=0.99,
+            xanchor="left",
+            x=0.01,
+            bgcolor="rgba(255,255,255,0.3)",
+        )
+    )
+    for trace in fig["data"]:
+        if trace["name"].startswith("Session"):
+            trace["showlegend"] = False
+    fig["data"][0]["name"] = "Individual Sessions"
+    fig["data"][0]["showlegend"] = True
 
     return fig
 
@@ -194,6 +209,7 @@ def _plot_all_benchmarks_combined_session_av(results_directory: Path) -> Figure:
         title=title,
     )
     fig["data"][0]["showlegend"] = True
+    fig.update_layout(legend=dict(yanchor="top", y=-0.2, xanchor="left", x=0.01, orientation="h"))
 
     return fig
 
@@ -248,14 +264,7 @@ def _plot_av_s_per_100_steps_10_nodes(
     versions = sorted(list(version_times_dict.keys()))
     times = [version_times_dict[version] for version in versions]
 
-    fig.add_trace(
+    fig.add_trace(go.Bar(x=versions, y=times, text=times, textposition="auto", texttemplate="%{y:.3f}"))
-        go.Bar(
-            x=versions,
-            y=times,
-            text=times,
-            textposition="auto",
-        )
-    )
 
     fig.update_layout(
         xaxis_title="PrimAITE Version",
@@ -267,7 +276,11 @@ def _plot_av_s_per_100_steps_10_nodes(
 
 
 def build_benchmark_md_report(
-    benchmark_start_time: datetime, session_metadata: Dict, config_path: Path, results_root_path: Path
+    benchmark_start_time: datetime,
+    session_metadata: Dict,
+    config_path: Path,
+    results_root_path: Path,
+    output_path: PathLike,
 ) -> None:
     """
     Generates a Markdown report for a benchmarking session, documenting performance metrics and graphs.
@@ -319,7 +332,7 @@ def build_benchmark_md_report(
     data = benchmark_metadata_dict
     primaite_version = data["primaite_version"]
 
-    with open(version_result_dir / f"PrimAITE v{primaite_version} Benchmark Report.md", "w") as file:
+    with open(output_path, "w") as file:
         # Title
         file.write(f"# PrimAITE v{primaite_version} Learning Benchmark\n")
         file.write("## PrimAITE Dev Team\n")
@@ -393,3 +406,15 @@ def build_benchmark_md_report(
             f"![Performance of Minor and Bugfix Releases for Major Version {major_v}]"
             f"({performance_benchmark_plot_path.name})\n"
         )
+
+
+def md2pdf(md_path: PathLike, pdf_path: PathLike, css_path: PathLike) -> None:
+    """Generate PDF version of Markdown report."""
+    from md2pdf.core import md2pdf
+
+    md2pdf(
+        pdf_file_path=pdf_path,
+        md_file_path=md_path,
+        base_url=Path(md_path).parent,
+        css_file_path=css_path,
+    )

benchmark/static/styles.css

@@ -0,0 +1,34 @@
+body {
+    font-family: 'Arial', sans-serif;
+    line-height: 1.6;
+    /* margin: 1cm; */
+}
+h1, h2, h3, h4, h5, h6 {
+    font-weight: bold;
+    /* margin: 1em 0; */
+}
+p {
+    /* margin: 0.5em 0; */
+}
+ul, ol {
+    margin: 1em 0;
+    padding-left: 1.5em;
+}
+pre {
+    background: #f4f4f4;
+    padding: 0.5em;
+    overflow-x: auto;
+}
+img {
+    max-width: 100%;
+    height: auto;
+}
+table {
+    width: 100%;
+    border-collapse: collapse;
+    margin: 1em 0;
+}
+th, td {
+    padding: 0.5em;
+    border: 1px solid #ddd;
+}

docs/source/simulation_components/system/common/common_configuration.rst

@@ -25,3 +25,35 @@ The configuration options are the attributes that fall under the options for an
 Optional. Default value is ``2``.
 
 The number of timesteps the |SOFTWARE_NAME| will remain in a ``FIXING`` state before going into a ``GOOD`` state.
+
+
+``listen_on_ports``
+"""""""""""""""""""
+
+The set of ports to listen on. This is in addition to the main port the software is designated. This set can either be
+the string name of ports or the port integers
+
+Example:
+
+.. code-block:: yaml
+
+    simulation:
+      network:
+        nodes:
+        - hostname: client
+          type: computer
+          ip_address: 192.168.10.11
+          subnet_mask: 255.255.255.0
+          default_gateway: 192.168.10.1
+          services:
+            - type: DatabaseService
+              options:
+                backup_server_ip: 10.10.1.12
+                listen_on_ports:
+                  - 631
+          applications:
+              - type: WebBrowser
+                options:
+                  target_url: http://sometech.ai
+                  listen_on_ports:
+                    - SMB

pyproject.toml

@@ -75,7 +75,8 @@ dev = [
     "wheel==0.38.4",
     "nbsphinx==0.9.4",
     "nbmake==1.5.4",
-    "pytest-xdist==3.3.1"
+    "pytest-xdist==3.3.1",
+    "md2pdf",
 ]
 
 [project.scripts]

src/primaite/game/game.py

@@ -1,7 +1,7 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 """PrimAITE game - Encapsulates the simulation and agents."""
 from ipaddress import IPv4Address
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 from pydantic import BaseModel, ConfigDict
@@ -44,8 +44,10 @@ from primaite.simulator.system.services.ftp.ftp_client import FTPClient
 from primaite.simulator.system.services.ftp.ftp_server import FTPServer
 from primaite.simulator.system.services.ntp.ntp_client import NTPClient
 from primaite.simulator.system.services.ntp.ntp_server import NTPServer
+from primaite.simulator.system.services.service import Service
 from primaite.simulator.system.services.terminal.terminal import Terminal
 from primaite.simulator.system.services.web_server.web_server import WebServer
+from primaite.simulator.system.software import Software
 
 _LOGGER = getLogger(__name__)
 
@@ -328,6 +330,20 @@ class PrimaiteGame:
                 user_manager: UserManager = new_node.software_manager.software["UserManager"]  # noqa
                 for user_cfg in node_cfg["users"]:
                     user_manager.add_user(**user_cfg, bypass_can_perform_action=True)
+
+            def _set_software_listen_on_ports(software: Union[Software, Service], software_cfg: Dict):
+                """Set listener ports on software."""
+                listen_on_ports = []
+                for port_id in set(software_cfg.get("options", {}).get("listen_on_ports", [])):
+                    port = None
+                    if isinstance(port_id, int):
+                        port = Port(port_id)
+                    elif isinstance(port_id, str):
+                        port = Port[port_id]
+                    if port:
+                        listen_on_ports.append(port)
+                software.listen_on_ports = set(listen_on_ports)
+
             if "services" in node_cfg:
                 for service_cfg in node_cfg["services"]:
                     new_service = None
@@ -341,6 +357,7 @@ class PrimaiteGame:
                         if "fix_duration" in service_cfg.get("options", {}):
                             new_service.fixing_duration = service_cfg["options"]["fix_duration"]
 
+                        _set_software_listen_on_ports(new_service, service_cfg)
                         # start the service
                         new_service.start()
                     else:
@@ -390,6 +407,8 @@ class PrimaiteGame:
                         _LOGGER.error(msg)
                         raise ValueError(msg)
 
+                    _set_software_listen_on_ports(new_application, application_cfg)
+
                     # run the application
                     new_application.run()
 

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

@@ -1,4 +1,5 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from copy import deepcopy
 from ipaddress import IPv4Address, IPv4Network
 from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
@@ -76,6 +77,8 @@ class SoftwareManager:
         for software in self.port_protocol_mapping.values():
             if software.operating_state in {ApplicationOperatingState.RUNNING, ServiceOperatingState.RUNNING}:
                 open_ports.append(software.port)
+                if software.listen_on_ports:
+                    open_ports += list(software.listen_on_ports)
         return open_ports
 
     def check_port_is_open(self, port: Port, protocol: IPProtocol) -> bool:
@@ -223,7 +226,9 @@ class SoftwareManager:
         frame: Frame,
     ):
         """
-        Receive a payload from the SessionManager and forward it to the corresponding service or application.
+        Receive a payload from the SessionManager and forward it to the corresponding service or applications.
+
+        This function handles both software assigned a specific port, and software listening in on other ports.
 
         :param payload: The payload being received.
         :param session: The transport session the payload originates from.
@@ -231,14 +236,25 @@ class SoftwareManager:
         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)
+        main_receiver = self.port_protocol_mapping.get((port, protocol), None)
-        if receiver:
+        if main_receiver:
-            receiver.receive(
+            main_receiver.receive(
                 payload=payload, session_id=session_id, from_network_interface=from_network_interface, frame=frame
             )
-        else:
+        listening_receivers = [
+            software
+            for software in self.software.values()
+            if port in software.listen_on_ports and software != main_receiver
+        ]
+        for receiver in listening_receivers:
+            receiver.receive(
+                payload=deepcopy(payload),
+                session_id=session_id,
+                from_network_interface=from_network_interface,
+                frame=frame,
+            )
+        if not main_receiver and not listening_receivers:
             self.sys_log.warning(f"No service or application found for port {port} and protocol {protocol}")
-        pass
 
     def show(self, markdown: bool = False):
         """

src/primaite/simulator/system/services/database/database_service.py

@@ -385,6 +385,8 @@ class DatabaseService(Service):
                     )
                 else:
                     result = {"status_code": 401, "type": "sql"}
+        else:
+            self.sys_log.info(f"{self.name}: Ignoring payload as it is not a Database payload")
         self.send(payload=result, session_id=session_id)
         return True
 

src/primaite/simulator/system/software.py

@@ -4,9 +4,10 @@ from abc import abstractmethod
 from datetime import datetime
 from enum import Enum
 from ipaddress import IPv4Address, IPv4Network
-from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Union
 
 from prettytable import MARKDOWN, PrettyTable
+from pydantic import Field
 
 from primaite.interface.request import RequestResponse
 from primaite.simulator.core import RequestManager, RequestType, SimComponent
@@ -252,6 +253,8 @@ class IOSoftware(Software):
     "Indicates if the software uses UDP protocol for communication. Default is True."
     port: Port
     "The port to which the software is connected."
+    listen_on_ports: Set[Port] = Field(default_factory=set)
+    "The set of ports to listen on."
     protocol: IPProtocol
     "The IP Protocol the Software operates on."
     _connections: Dict[str, Dict] = {}

tests/assets/configs/basic_node_with_software_listening_ports.yaml

@@ -0,0 +1,39 @@
+io_settings:
+  save_step_metadata: false
+  save_pcap_logs: true
+  save_sys_logs: true
+  sys_log_level: WARNING
+  agent_log_level: INFO
+  save_agent_logs: true
+  write_agent_log_to_terminal: True
+
+
+game:
+  max_episode_length: 256
+  ports:
+  - ARP
+  protocols:
+  - ICMP
+  - UDP
+
+
+simulation:
+  network:
+    nodes:
+    - hostname: client
+      type: computer
+      ip_address: 192.168.10.11
+      subnet_mask: 255.255.255.0
+      default_gateway: 192.168.10.1
+      services:
+        - type: DatabaseService
+          options:
+            backup_server_ip: 10.10.1.12
+            listen_on_ports:
+              - 631
+      applications:
+          - type: WebBrowser
+            options:
+              target_url: http://sometech.ai
+              listen_on_ports:
+                - SMB

tests/integration_tests/system/test_service_listening_on_ports.py

@@ -0,0 +1,84 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from typing import Any, Dict, List, Set
+
+import yaml
+from pydantic import Field
+
+from primaite.game.game import PrimaiteGame
+from primaite.simulator.network.hardware.nodes.host.computer import Computer
+from primaite.simulator.network.transmission.network_layer import IPProtocol
+from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.applications.database_client import DatabaseClient
+from primaite.simulator.system.services.database.database_service import DatabaseService
+from primaite.simulator.system.services.service import Service
+from tests import TEST_ASSETS_ROOT
+
+
+class _DatabaseListener(Service):
+    name: str = "DatabaseListener"
+    protocol: IPProtocol = IPProtocol.TCP
+    port: Port = Port.NONE
+    listen_on_ports: Set[Port] = {Port.POSTGRES_SERVER}
+    payloads_received: List[Any] = Field(default_factory=list)
+
+    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
+        self.payloads_received.append(payload)
+        self.sys_log.info(f"{self.name}: received payload {payload}")
+        return True
+
+    def describe_state(self) -> Dict:
+        return super().describe_state()
+
+
+def test_http_listener(client_server):
+    computer, server = client_server
+
+    server.software_manager.install(DatabaseService)
+    server_db = server.software_manager.software["DatabaseService"]
+    server_db.start()
+
+    server.software_manager.install(_DatabaseListener)
+    server_db_listener: _DatabaseListener = server.software_manager.software["DatabaseListener"]
+    server_db_listener.start()
+
+    computer.software_manager.install(DatabaseClient)
+    computer_db_client: DatabaseClient = computer.software_manager.software["DatabaseClient"]
+
+    computer_db_client.run()
+    computer_db_client.server_ip_address = server.network_interface[1].ip_address
+
+    assert len(server_db_listener.payloads_received) == 0
+    computer.session_manager.receive_payload_from_software_manager(
+        payload="masquerade as Database traffic",
+        dst_ip_address=server.network_interface[1].ip_address,
+        dst_port=Port.POSTGRES_SERVER,
+        ip_protocol=IPProtocol.TCP,
+    )
+
+    assert len(server_db_listener.payloads_received) == 1
+
+    db_connection = computer_db_client.get_new_connection()
+
+    assert db_connection
+
+    assert len(server_db_listener.payloads_received) == 2
+
+    assert db_connection.query("SELECT")
+
+    assert len(server_db_listener.payloads_received) == 3
+
+
+def test_set_listen_on_ports_from_config():
+    config_path = TEST_ASSETS_ROOT / "configs" / "basic_node_with_software_listening_ports.yaml"
+
+    with open(config_path, "r") as f:
+        config_dict = yaml.safe_load(f)
+    network = PrimaiteGame.from_config(cfg=config_dict).simulation.network
+
+    client: Computer = network.get_node_by_hostname("client")
+    assert Port.SMB in client.software_manager.get_open_ports()
+    assert Port.IPP in client.software_manager.get_open_ports()
+
+    web_browser = client.software_manager.software["WebBrowser"]
+
+    assert not web_browser.listen_on_ports.difference({Port.SMB, Port.IPP})