#2725: add fix duration to application and service configuration

src/primaite/game/game.py

@@ -297,6 +297,11 @@ class PrimaiteGame:
                         new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
                         new_service = new_node.software_manager.software[service_type]
 
+                        # fixing duration for the service
+                        fix_duration = service_cfg.get("options", {}).get("fix_duration", None)
+                        if fix_duration:
+                            new_service.fixing_duration = fix_duration
+
                         # start the service
                         new_service.start()
                     else:
@@ -336,6 +341,11 @@ class PrimaiteGame:
                     if application_type in APPLICATION_TYPES_MAPPING:
                         new_node.software_manager.install(APPLICATION_TYPES_MAPPING[application_type])
                         new_application = new_node.software_manager.software[application_type]
+
+                        # fixing duration for the application
+                        fix_duration = application_cfg.get("options", {}).get("fix_duration", None)
+                        if fix_duration:
+                            new_application.fixing_duration = fix_duration
                     else:
                         msg = f"Configuration contains an invalid application type: {application_type}"
                         _LOGGER.error(msg)
@@ -358,7 +368,7 @@ class PrimaiteGame:
                         if "options" in application_cfg:
                             opt = application_cfg["options"]
                             new_application.configure(
-                                server_ip_address=IPv4Address(opt.get("server_ip")),
+                                server_ip_address=IPv4Address(opt.get("server_ip")) if opt.get("server_ip") else None,
                                 server_password=opt.get("server_password"),
                                 payload=opt.get("payload", "ENCRYPT"),
                             )

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

@@ -88,7 +88,7 @@ class RansomwareScript(Application):
 
     def configure(
         self,
-        server_ip_address: IPv4Address,
+        server_ip_address: Optional[IPv4Address] = None,
         server_password: Optional[str] = None,
         payload: Optional[str] = None,
     ):

tests/assets/configs/fix_duration_one_item.yaml

@@ -0,0 +1,248 @@
+#  Basic Switched network
+#
+#  --------------      --------------      --------------
+#  |  client_1  |------|  switch_1  |------|  client_2  |
+#  --------------      --------------      --------------
+#
+io_settings:
+  save_step_metadata: false
+  save_pcap_logs: true
+  save_sys_logs: true
+  sys_log_level: WARNING
+
+
+game:
+  max_episode_length: 256
+  ports:
+  - ARP
+  - DNS
+  - HTTP
+  - POSTGRES_SERVER
+  protocols:
+  - ICMP
+  - TCP
+  - UDP
+
+agents:
+  - ref: client_2_green_user
+    team: GREEN
+    type: ProbabilisticAgent
+    observation_space: null
+    action_space:
+      action_list:
+        - type: DONOTHING
+        - type: NODE_APPLICATION_EXECUTE
+      action_map:
+        0:
+          action: DONOTHING
+          options: {}
+        1:
+          action: NODE_APPLICATION_EXECUTE
+          options:
+            node_id: 0
+            application_id: 0
+      options:
+        nodes:
+        - node_name: client_2
+          applications:
+            - application_name: WebBrowser
+        max_folders_per_node: 1
+        max_files_per_folder: 1
+        max_services_per_node: 1
+        max_applications_per_node: 1
+
+    reward_function:
+      reward_components:
+        - type: DUMMY
+
+    agent_settings:
+      start_settings:
+        start_step: 5
+        frequency: 4
+        variance: 3
+
+
+
+  - ref: defender
+    team: BLUE
+    type: ProxyAgent
+
+    observation_space:
+      type: CUSTOM
+      options:
+        components:
+          - type: NODES
+            label: NODES
+            options:
+              hosts:
+                - hostname: client_1
+                - hostname: client_2
+                - hostname: client_3
+              num_services: 1
+              num_applications: 0
+              num_folders: 1
+              num_files: 1
+              num_nics: 2
+              include_num_access: false
+              monitored_traffic:
+                icmp:
+                    - NONE
+                tcp:
+                    - DNS
+              include_nmne: true
+              routers:
+                - hostname: router_1
+              num_ports: 0
+              ip_list:
+                - 192.168.10.21
+                - 192.168.10.22
+                - 192.168.10.23
+              wildcard_list:
+                - 0.0.0.1
+              port_list:
+                - 80
+                - 5432
+              protocol_list:
+                - ICMP
+                - TCP
+                - UDP
+              num_rules: 10
+
+          - type: LINKS
+            label: LINKS
+            options:
+              link_references:
+                - switch_1:eth-1<->client_1:eth-1
+                - switch_1:eth-2<->client_2:eth-1
+          - type: "NONE"
+            label: ICS
+            options: {}
+
+    action_space:
+      action_list:
+        - type: DONOTHING
+
+      action_map:
+          0:
+            action: DONOTHING
+            options: {}
+      options:
+        nodes:
+        - node_name: switch
+        - node_name: client_1
+        - node_name: client_2
+        - node_name: client_3
+        max_folders_per_node: 2
+        max_files_per_folder: 2
+        max_services_per_node: 2
+        max_nics_per_node: 8
+        max_acl_rules: 10
+        ip_list:
+          - 192.168.10.21
+          - 192.168.10.22
+          - 192.168.10.23
+
+    reward_function:
+      reward_components:
+        - type: DATABASE_FILE_INTEGRITY
+          weight: 0.5
+          options:
+            node_hostname: database_server
+            folder_name: database
+            file_name: database.db
+
+
+        - type: WEB_SERVER_404_PENALTY
+          weight: 0.5
+          options:
+            node_hostname: web_server
+            service_name: web_server_web_service
+
+
+    agent_settings:
+      flatten_obs: true
+
+simulation:
+  network:
+    nodes:
+
+    - type: switch
+      hostname: switch_1
+      num_ports: 8
+
+    - hostname: client_1
+      type: computer
+      ip_address: 192.168.10.21
+      subnet_mask: 255.255.255.0
+      default_gateway: 192.168.10.1
+      dns_server: 192.168.1.10
+      applications:
+      - type: RansomwareScript
+      - type: WebBrowser
+        options:
+          target_url: http://arcd.com/users/
+      - type: DatabaseClient
+        options:
+          db_server_ip: 192.168.1.10
+          server_password: arcd
+          fix_duration: 1
+      - type: DataManipulationBot
+        options:
+          port_scan_p_of_success: 0.8
+          data_manipulation_p_of_success: 0.8
+          payload: "DELETE"
+          server_ip: 192.168.1.21
+          server_password: arcd
+      - type: DoSBot
+        options:
+          target_ip_address: 192.168.10.21
+          payload: SPOOF DATA
+          port_scan_p_of_success: 0.8
+      services:
+      - type: DNSClient
+        options:
+          dns_server: 192.168.1.10
+      - type: DNSServer
+        options:
+          domain_mapping:
+            arcd.com: 192.168.1.10
+      - type: DatabaseService
+        options:
+          fix_duration: 5
+          backup_server_ip: 192.168.1.10
+      - type: WebServer
+      - type: FTPClient
+      - type: FTPServer
+        options:
+          server_password: arcd
+      - type: NTPClient
+        options:
+          ntp_server_ip: 192.168.1.10
+      - type: NTPServer
+    - hostname: client_2
+      type: computer
+      ip_address: 192.168.10.22
+      subnet_mask: 255.255.255.0
+      default_gateway: 192.168.10.1
+      dns_server: 192.168.1.10
+      applications:
+        - type: DatabaseClient
+          options:
+            db_server_ip: 192.168.1.10
+            server_password: arcd
+      services:
+        - type: DNSClient
+          options:
+            dns_server: 192.168.1.10
+
+    links:
+    - endpoint_a_hostname: switch_1
+      endpoint_a_port: 1
+      endpoint_b_hostname: client_1
+      endpoint_b_port: 1
+      bandwidth: 200
+    - endpoint_a_hostname: switch_1
+      endpoint_a_port: 2
+      endpoint_b_hostname: client_2
+      endpoint_b_port: 1
+      bandwidth: 200

tests/assets/configs/software_fix_duration.yaml

@@ -0,0 +1,263 @@
+#  Basic Switched network
+#
+#  --------------      --------------      --------------
+#  |  client_1  |------|  switch_1  |------|  client_2  |
+#  --------------      --------------      --------------
+#
+io_settings:
+  save_step_metadata: false
+  save_pcap_logs: true
+  save_sys_logs: true
+  sys_log_level: WARNING
+
+
+game:
+  max_episode_length: 256
+  ports:
+  - ARP
+  - DNS
+  - HTTP
+  - POSTGRES_SERVER
+  protocols:
+  - ICMP
+  - TCP
+  - UDP
+
+agents:
+  - ref: client_2_green_user
+    team: GREEN
+    type: ProbabilisticAgent
+    observation_space: null
+    action_space:
+      action_list:
+        - type: DONOTHING
+        - type: NODE_APPLICATION_EXECUTE
+      action_map:
+        0:
+          action: DONOTHING
+          options: {}
+        1:
+          action: NODE_APPLICATION_EXECUTE
+          options:
+            node_id: 0
+            application_id: 0
+      options:
+        nodes:
+        - node_name: client_2
+          applications:
+            - application_name: WebBrowser
+        max_folders_per_node: 1
+        max_files_per_folder: 1
+        max_services_per_node: 1
+        max_applications_per_node: 1
+
+    reward_function:
+      reward_components:
+        - type: DUMMY
+
+    agent_settings:
+      start_settings:
+        start_step: 5
+        frequency: 4
+        variance: 3
+
+
+
+  - ref: defender
+    team: BLUE
+    type: ProxyAgent
+
+    observation_space:
+      type: CUSTOM
+      options:
+        components:
+          - type: NODES
+            label: NODES
+            options:
+              hosts:
+                - hostname: client_1
+                - hostname: client_2
+                - hostname: client_3
+              num_services: 1
+              num_applications: 0
+              num_folders: 1
+              num_files: 1
+              num_nics: 2
+              include_num_access: false
+              monitored_traffic:
+                icmp:
+                    - NONE
+                tcp:
+                    - DNS
+              include_nmne: true
+              routers:
+                - hostname: router_1
+              num_ports: 0
+              ip_list:
+                - 192.168.10.21
+                - 192.168.10.22
+                - 192.168.10.23
+              wildcard_list:
+                - 0.0.0.1
+              port_list:
+                - 80
+                - 5432
+              protocol_list:
+                - ICMP
+                - TCP
+                - UDP
+              num_rules: 10
+
+          - type: LINKS
+            label: LINKS
+            options:
+              link_references:
+                - switch_1:eth-1<->client_1:eth-1
+                - switch_1:eth-2<->client_2:eth-1
+          - type: "NONE"
+            label: ICS
+            options: {}
+
+    action_space:
+      action_list:
+        - type: DONOTHING
+
+      action_map:
+          0:
+            action: DONOTHING
+            options: {}
+      options:
+        nodes:
+        - node_name: switch
+        - node_name: client_1
+        - node_name: client_2
+        - node_name: client_3
+        max_folders_per_node: 2
+        max_files_per_folder: 2
+        max_services_per_node: 2
+        max_nics_per_node: 8
+        max_acl_rules: 10
+        ip_list:
+          - 192.168.10.21
+          - 192.168.10.22
+          - 192.168.10.23
+
+    reward_function:
+      reward_components:
+        - type: DATABASE_FILE_INTEGRITY
+          weight: 0.5
+          options:
+            node_hostname: database_server
+            folder_name: database
+            file_name: database.db
+
+
+        - type: WEB_SERVER_404_PENALTY
+          weight: 0.5
+          options:
+            node_hostname: web_server
+            service_name: web_server_web_service
+
+
+    agent_settings:
+      flatten_obs: true
+
+simulation:
+  network:
+    nodes:
+
+    - type: switch
+      hostname: switch_1
+      num_ports: 8
+
+    - hostname: client_1
+      type: computer
+      ip_address: 192.168.10.21
+      subnet_mask: 255.255.255.0
+      default_gateway: 192.168.10.1
+      dns_server: 192.168.1.10
+      applications:
+      - type: RansomwareScript
+        options:
+          fix_duration: 1
+      - type: WebBrowser
+        options:
+          target_url: http://arcd.com/users/
+          fix_duration: 1
+      - type: DatabaseClient
+        options:
+          db_server_ip: 192.168.1.10
+          server_password: arcd
+          fix_duration: 1
+      - type: DataManipulationBot
+        options:
+          port_scan_p_of_success: 0.8
+          data_manipulation_p_of_success: 0.8
+          payload: "DELETE"
+          server_ip: 192.168.1.21
+          server_password: arcd
+          fix_duration: 1
+      - type: DoSBot
+        options:
+          target_ip_address: 192.168.10.21
+          payload: SPOOF DATA
+          port_scan_p_of_success: 0.8
+          fix_duration: 1
+      services:
+      - type: DNSClient
+        options:
+          dns_server: 192.168.1.10
+          fix_duration: 3
+      - type: DNSServer
+        options:
+          fix_duration: 3
+          domain_mapping:
+            arcd.com: 192.168.1.10
+      - type: DatabaseService
+        options:
+          backup_server_ip: 192.168.1.10
+          fix_duration: 3
+      - type: WebServer
+        options:
+          fix_duration: 3
+      - type: FTPClient
+        options:
+          fix_duration: 3
+      - type: FTPServer
+        options:
+          server_password: arcd
+          fix_duration: 3
+      - type: NTPClient
+        options:
+          ntp_server_ip: 192.168.1.10
+          fix_duration: 3
+      - type: NTPServer
+        options:
+          fix_duration: 3
+    - hostname: client_2
+      type: computer
+      ip_address: 192.168.10.22
+      subnet_mask: 255.255.255.0
+      default_gateway: 192.168.10.1
+      dns_server: 192.168.1.10
+      applications:
+        - type: DatabaseClient
+          options:
+            db_server_ip: 192.168.1.10
+            server_password: arcd
+      services:
+        - type: DNSClient
+          options:
+            dns_server: 192.168.1.10
+
+    links:
+    - endpoint_a_hostname: switch_1
+      endpoint_a_port: 1
+      endpoint_b_hostname: client_1
+      endpoint_b_port: 1
+      bandwidth: 200
+    - endpoint_a_hostname: switch_1
+      endpoint_a_port: 2
+      endpoint_b_hostname: client_2
+      endpoint_b_port: 1
+      bandwidth: 200

tests/integration_tests/configuration_file_parsing/test_software_fix_duration.py

@@ -0,0 +1,93 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+import copy
+from ipaddress import IPv4Address
+from pathlib import Path
+from typing import Union
+
+import yaml
+
+from primaite.config.load import data_manipulation_config_path
+from primaite.game.agent.interface import ProxyAgent
+from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
+from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
+from primaite.game.game import APPLICATION_TYPES_MAPPING, PrimaiteGame, SERVICE_TYPES_MAPPING
+from primaite.simulator.network.container import Network
+from primaite.simulator.network.hardware.nodes.host.computer import Computer
+from primaite.simulator.system.applications.database_client import DatabaseClient
+from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
+from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
+from primaite.simulator.system.applications.web_browser import WebBrowser
+from primaite.simulator.system.services.database.database_service import DatabaseService
+from primaite.simulator.system.services.dns.dns_client import DNSClient
+from primaite.simulator.system.services.dns.dns_server import DNSServer
+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.web_server.web_server import WebServer
+from tests import TEST_ASSETS_ROOT
+
+TEST_CONFIG = TEST_ASSETS_ROOT / "configs/software_fix_duration.yaml"
+ONE_ITEM_CONFIG = TEST_ASSETS_ROOT / "configs/fix_duration_one_item.yaml"
+
+
+def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
+    """Returns a PrimaiteGame object which loads the contents of a given yaml path."""
+    with open(config_path, "r") as f:
+        cfg = yaml.safe_load(f)
+
+    return PrimaiteGame.from_config(cfg)
+
+
+def test_default_fix_duration():
+    """Test that software with no defined fix duration in config uses the default fix duration of 2."""
+    game = load_config(TEST_CONFIG)
+    client_2: Computer = game.simulation.network.get_node_by_hostname("client_2")
+
+    database_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient")
+    assert database_client.fixing_duration == 2
+
+    dns_client: DNSClient = client_2.software_manager.software.get("DNSClient")
+    assert dns_client.fixing_duration == 2
+
+
+def test_fix_duration_set_from_config():
+    """Test to check that the fix duration set for applications and services works as intended."""
+    game = load_config(TEST_CONFIG)
+    client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
+
+    # in config - services take 3 timesteps to fix
+    for service in SERVICE_TYPES_MAPPING:
+        assert client_1.software_manager.software.get(service) is not None
+        assert client_1.software_manager.software.get(service).fixing_duration == 3
+
+    # in config - applications take 1 timestep to fix
+    for applications in APPLICATION_TYPES_MAPPING:
+        assert client_1.software_manager.software.get(applications) is not None
+        assert client_1.software_manager.software.get(applications).fixing_duration == 1
+
+
+def test_fix_duration_for_one_item():
+    """Test that setting fix duration for one application does not affect other components."""
+    game = load_config(ONE_ITEM_CONFIG)
+    client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
+
+    # in config - services take 3 timesteps to fix
+    services = copy.copy(SERVICE_TYPES_MAPPING)
+    services.pop("DatabaseService")
+    for service in services:
+        assert client_1.software_manager.software.get(service) is not None
+        assert client_1.software_manager.software.get(service).fixing_duration == 2
+
+    # in config - applications take 1 timestep to fix
+    applications = copy.copy(APPLICATION_TYPES_MAPPING)
+    applications.pop("DatabaseClient")
+    for applications in applications:
+        assert client_1.software_manager.software.get(applications) is not None
+        assert client_1.software_manager.software.get(applications).fixing_duration == 2
+
+    database_client: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
+    assert database_client.fixing_duration == 1
+
+    database_service: DatabaseService = client_1.software_manager.software.get("DatabaseService")
+    assert database_service.fixing_duration == 5