From 7c26ca9d79d14bc529368f455909a097c8a4614c Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 12 Sep 2024 16:07:14 +0100 Subject: [PATCH] #2864: add configuration for services_requires_scan and applications_requires_scan --- .../agent/observations/host_observations.py | 24 +++- .../agent/observations/node_observations.py | 12 +- .../observations/software_observation.py | 33 ++++-- .../_game/_agent/test_observations.py | 111 +++++++++++++++++- 4 files changed, 169 insertions(+), 11 deletions(-) diff --git a/src/primaite/game/agent/observations/host_observations.py b/src/primaite/game/agent/observations/host_observations.py index 4419ccc7..3371a99c 100644 --- a/src/primaite/game/agent/observations/host_observations.py +++ b/src/primaite/game/agent/observations/host_observations.py @@ -52,6 +52,14 @@ class HostObservation(AbstractObservation, identifier="HOST"): """ If True, files and folders must be scanned to update the health state. If False, true state is always shown. """ + services_requires_scan: Optional[bool] = None + """ + If True, services must be scanned to update the health state. If False, true state is always shown. + """ + applications_requires_scan: Optional[bool] = None + """ + If True, applications must be scanned to update the health state. If False, true state is always shown. + """ include_users: Optional[bool] = True """If True, report user session information.""" @@ -71,6 +79,8 @@ class HostObservation(AbstractObservation, identifier="HOST"): monitored_traffic: Optional[Dict], include_num_access: bool, file_system_requires_scan: bool, + services_requires_scan: bool, + applications_requires_scan: bool, include_users: bool, ) -> None: """ @@ -106,6 +116,12 @@ class HostObservation(AbstractObservation, identifier="HOST"): :param file_system_requires_scan: If True, the files and folders must be scanned to update the health state. If False, the true state is always shown. :type file_system_requires_scan: bool + :param services_requires_scan: If True, services must be scanned to update the health state. + If False, the true state is always shown. + :type services_requires_scan: bool + :param applications_requires_scan: If True, applications must be scanned to update the health state. + If False, the true state is always shown. + :type applications_requires_scan: bool :param include_users: If True, report user session information. :type include_users: bool """ @@ -119,7 +135,7 @@ class HostObservation(AbstractObservation, identifier="HOST"): # Ensure lists have lengths equal to specified counts by truncating or padding self.services: List[ServiceObservation] = services while len(self.services) < num_services: - self.services.append(ServiceObservation(where=None)) + self.services.append(ServiceObservation(where=None, services_requires_scan=services_requires_scan)) while len(self.services) > num_services: truncated_service = self.services.pop() msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}" @@ -127,7 +143,9 @@ class HostObservation(AbstractObservation, identifier="HOST"): self.applications: List[ApplicationObservation] = applications while len(self.applications) < num_applications: - self.applications.append(ApplicationObservation(where=None)) + self.applications.append( + ApplicationObservation(where=None, applications_requires_scan=applications_requires_scan) + ) while len(self.applications) > num_applications: truncated_application = self.applications.pop() msg = f"Too many applications in Node observation space for node. Truncating {truncated_application.where}" @@ -293,5 +311,7 @@ class HostObservation(AbstractObservation, identifier="HOST"): monitored_traffic=config.monitored_traffic, include_num_access=config.include_num_access, file_system_requires_scan=config.file_system_requires_scan, + services_requires_scan=config.services_requires_scan, + applications_requires_scan=config.applications_requires_scan, include_users=config.include_users, ) diff --git a/src/primaite/game/agent/observations/node_observations.py b/src/primaite/game/agent/observations/node_observations.py index e263cadb..85de5396 100644 --- a/src/primaite/game/agent/observations/node_observations.py +++ b/src/primaite/game/agent/observations/node_observations.py @@ -45,7 +45,13 @@ class NodesObservation(AbstractObservation, identifier="NODES"): include_num_access: Optional[bool] = None """Flag to include the number of accesses.""" file_system_requires_scan: bool = True - """If True, the folder must be scanned to update the health state. Tf False, the true state is always shown.""" + """If True, the folder must be scanned to update the health state. If False, the true state is always shown.""" + services_requires_scan: bool = True + """If True, the services must be scanned to update the health state. + If False, the true state is always shown.""" + applications_requires_scan: bool = True + """If True, the applications must be scanned to update the health state. + If False, the true state is always shown.""" include_users: Optional[bool] = True """If True, report user session information.""" num_ports: Optional[int] = None @@ -193,6 +199,10 @@ class NodesObservation(AbstractObservation, identifier="NODES"): host_config.include_num_access = config.include_num_access if host_config.file_system_requires_scan is None: host_config.file_system_requires_scan = config.file_system_requires_scan + if host_config.services_requires_scan is None: + host_config.services_requires_scan = config.services_requires_scan + if host_config.applications_requires_scan is None: + host_config.applications_requires_scan = config.applications_requires_scan if host_config.include_users is None: host_config.include_users = config.include_users diff --git a/src/primaite/game/agent/observations/software_observation.py b/src/primaite/game/agent/observations/software_observation.py index 15cd2447..2075ce43 100644 --- a/src/primaite/game/agent/observations/software_observation.py +++ b/src/primaite/game/agent/observations/software_observation.py @@ -1,7 +1,7 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from __future__ import annotations -from typing import Dict +from typing import Dict, Optional from gymnasium import spaces from gymnasium.core import ObsType @@ -19,7 +19,10 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"): service_name: str """Name of the service, used for querying simulation state dictionary""" - def __init__(self, where: WhereType) -> None: + services_requires_scan: Optional[bool] = None + """If True, services must be scanned to update the health state. If False, true state is always shown.""" + + def __init__(self, where: WhereType, services_requires_scan: bool) -> None: """ Initialise a service observation instance. @@ -28,6 +31,7 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"): :type where: WhereType """ self.where = where + self.services_requires_scan = services_requires_scan self.default_observation = {"operating_status": 0, "health_status": 0} def observe(self, state: Dict) -> ObsType: @@ -44,7 +48,9 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"): return self.default_observation return { "operating_status": service_state["operating_state"], - "health_status": service_state["health_state_visible"], + "health_status": service_state["health_state_visible"] + if self.services_requires_scan + else service_state["health_state_actual"], } @property @@ -70,7 +76,9 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"): :return: Constructed service observation instance. :rtype: ServiceObservation """ - return cls(where=parent_where + ["services", config.service_name]) + return cls( + where=parent_where + ["services", config.service_name], services_requires_scan=config.services_requires_scan + ) class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): @@ -82,7 +90,12 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): application_name: str """Name of the application, used for querying simulation state dictionary""" - def __init__(self, where: WhereType) -> None: + applications_requires_scan: Optional[bool] = None + """ + If True, applications must be scanned to update the health state. If False, true state is always shown. + """ + + def __init__(self, where: WhereType, applications_requires_scan: bool) -> None: """ Initialise an application observation instance. @@ -92,6 +105,7 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): :type where: WhereType """ self.where = where + self.applications_requires_scan = applications_requires_scan self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0} # TODO: allow these to be configured in yaml @@ -128,7 +142,9 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): return self.default_observation return { "operating_status": application_state["operating_state"], - "health_status": application_state["health_state_visible"], + "health_status": application_state["health_state_visible"] + if self.applications_requires_scan + else application_state["health_state_actual"], "num_executions": self._categorise_num_executions(application_state["num_executions"]), } @@ -161,4 +177,7 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): :return: Constructed application observation instance. :rtype: ApplicationObservation """ - return cls(where=parent_where + ["applications", config.application_name]) + return cls( + where=parent_where + ["applications", config.application_name], + applications_requires_scan=config.applications_requires_scan, + ) diff --git a/tests/unit_tests/_primaite/_game/_agent/test_observations.py b/tests/unit_tests/_primaite/_game/_agent/test_observations.py index 7f590685..583b9cbd 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_observations.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_observations.py @@ -4,7 +4,7 @@ from typing import List import pytest import yaml -from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.observations import ObservationManager, ServiceObservation from primaite.game.agent.observations.file_system_observations import FileObservation, FolderObservation from primaite.game.agent.observations.host_observations import HostObservation @@ -130,3 +130,112 @@ class TestFileSystemRequiresScan: [], files=[], num_files=0, include_num_access=False, file_system_requires_scan=False ) assert obs_not_requiring_scan.observe(folder_state)["health_status"] == 3 + + +class TestServiceRequiresScan: + @pytest.mark.parametrize( + ("yaml_option_string", "expected_val"), + ( + ("services_requires_scan: true", True), + ("services_requires_scan: false", False), + (" ", True), + ), + ) + def test_obs_config(self, yaml_option_string, expected_val): + """Check that the default behaviour is to set FileSystemRequiresScan to True.""" + obs_cfg_yaml = f""" + type: CUSTOM + options: + components: + - type: NODES + label: NODES + options: + hosts: + - hostname: domain_controller + - hostname: web_server + services: + - service_name: WebServer + - hostname: database_server + folders: + - folder_name: database + files: + - file_name: database.db + - hostname: backup_server + - hostname: security_suite + - hostname: client_1 + applications: + - application_name: WebBrowser + - hostname: client_2 + num_services: 1 + num_applications: 1 + num_folders: 1 + num_files: 1 + num_nics: 2 + include_num_access: false + {yaml_option_string} + include_nmne: true + monitored_traffic: + icmp: + - NONE + tcp: + - DNS + routers: + - hostname: router_1 + num_ports: 0 + ip_list: + - 192.168.1.10 + - 192.168.1.12 + - 192.168.1.14 + - 192.168.1.16 + - 192.168.1.110 + - 192.168.10.21 + - 192.168.10.22 + - 192.168.10.110 + 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: + - router_1:eth-1<->switch_1:eth-8 + - router_1:eth-2<->switch_2:eth-8 + - switch_1:eth-1<->domain_controller:eth-1 + - switch_1:eth-2<->web_server:eth-1 + - switch_1:eth-3<->database_server:eth-1 + - switch_1:eth-4<->backup_server:eth-1 + - switch_1:eth-7<->security_suite:eth-1 + - switch_2:eth-1<->client_1:eth-1 + - switch_2:eth-2<->client_2:eth-1 + - switch_2:eth-7<->security_suite:eth-2 + - type: "NONE" + label: ICS + options: {{}} + + """ + + cfg = yaml.safe_load(obs_cfg_yaml) + manager = ObservationManager.from_config(cfg) + + hosts: List[HostObservation] = manager.obs.components["NODES"].hosts + for host in hosts: + services: List[ServiceObservation] = host.services + for service in services: + assert service.services_requires_scan == expected_val # Make sure services require scan by default + + def test_services_requires_scan(self): + state = {"health_state_actual": 3, "health_state_visible": 1, "operating_state": 1} + + obs_requiring_scan = ServiceObservation([], services_requires_scan=True) + assert obs_requiring_scan.observe(state)["health_status"] == 1 # should be visible value + + obs_not_requiring_scan = ServiceObservation([], services_requires_scan=False) + assert obs_not_requiring_scan.observe(state)["health_status"] == 3 # should be actual value