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

#2456 - Merging in updates on Dev and resolving merge conflicts

Browse Source
This commit is contained in:
Charlie Crane
2024-09-13 11:55:26 +01:00
parent 17035be028 519fb98e25
commit 59e1ca6615
11 changed files with 525 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/primaite/game/agent/actions.py
View File

@@ -1116,6 +1116,38 @@ class ConfigureC2BeaconAction(AbstractAction):
return ["network", "node", node_name, "application", "C2Beacon", "configure", config.__dict__]
class NodeAccountsAddUserAction(AbstractAction):
"""Action which changes adds a User."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: str, username: str, password: str, is_admin: bool) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
return ["network", "node", node_name, "service", "UserManager", "add_user", username, password, is_admin]
class NodeAccountsDisableUserAction(AbstractAction):
"""Action which disables a user."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: str, username: str) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
return [
"network",
"node",
node_name,
"service",
"UserManager",
"disable_user",
username,
]
class NodeAccountsChangePasswordAction(AbstractAction):
"""Action which changes the password for a user."""
@@ -1368,6 +1400,8 @@ class ActionManager:
"C2_SERVER_RANSOMWARE_CONFIGURE": RansomwareConfigureC2ServerAction,
"C2_SERVER_TERMINAL_COMMAND": TerminalC2ServerAction,
"C2_SERVER_DATA_EXFILTRATE": ExfiltrationC2ServerAction,
"NODE_ACCOUNTS_ADD_USER": NodeAccountsAddUserAction,
"NODE_ACCOUNTS_DISABLE_USER": NodeAccountsDisableUserAction,
"NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction,
"SSH_TO_REMOTE": NodeSessionsRemoteLoginAction,
"SESSIONS_REMOTE_LOGOFF": NodeSessionsRemoteLogoutAction,

28
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}"
@@ -263,6 +281,10 @@ class HostObservation(AbstractObservation, identifier="HOST"):
folder_config.file_system_requires_scan = config.file_system_requires_scan
for nic_config in config.network_interfaces:
nic_config.include_nmne = config.include_nmne
for service_config in config.services:
service_config.services_requires_scan = config.services_requires_scan
for application_config in config.applications:
application_config.applications_requires_scan = config.applications_requires_scan
services = [ServiceObservation.from_config(config=c, parent_where=where) for c in config.services]
applications = [ApplicationObservation.from_config(config=c, parent_where=where) for c in config.applications]
@@ -293,5 +315,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,
)

16
src/primaite/simulator/network/hardware/base.py
View File

@@ -857,7 +857,21 @@ class UserManager(Service):
"""
rm = super()._init_request_manager()
# todo add doc about requeest schemas
# todo add doc about request schemas
rm.add_request(
"add_user",
RequestType(
func=lambda request, context: RequestResponse.from_bool(
self.add_user(username=request[0], password=request[1], is_admin=request[2])
)
),
)
rm.add_request(
"disable_user",
RequestType(
func=lambda request, context: RequestResponse.from_bool(self.disable_user(username=request[0]))
),
)
rm.add_request(
"change_password",
RequestType(