oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2864: add configuration for services_requires_scan and applications_requires_scan

Browse Source
This commit is contained in:
Czar Echavez
2024-09-12 16:07:14 +01:00
parent 31a820f021
commit 7c26ca9d79
4 changed files with 169 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/primaite/game/agent/observations/host_observations.py
View File

@@ -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,
)

12
src/primaite/game/agent/observations/node_observations.py
View File

@@ -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

33
src/primaite/game/agent/observations/software_observation.py
View File

@@ -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,
)

111
tests/unit_tests/_primaite/_game/_agent/test_observations.py
View File

@@ -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