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

#3029 - Add string-ip validator, improve validation, fix minor bugs in pulling schema data

Browse Source
This commit is contained in:
Marek Wolan
2025-01-31 12:18:52 +00:00
parent 09bfa7b700
commit 4fb54c9492
19 changed files with 116 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/primaite/game/agent/actions/software.py
View File

@@ -6,6 +6,9 @@ from pydantic import ConfigDict, Field
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.ipv4_address import StrIP
from primaite.utils.validation.port import Port
__all__ = (
"ConfigureRansomwareScriptAction",
@@ -64,8 +67,8 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
model_config = ConfigDict(extra="forbid")
node_name: str
target_ip_address: Optional[str] = None
target_port: Optional[str] = None
target_ip_address: Optional[StrIP] = None
target_port: Optional[Port] = None
payload: Optional[str] = None
repeat: Optional[bool] = None
port_scan_p_of_success: Optional[float] = None
@@ -95,10 +98,10 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
"""Configuration schema for ConfigureC2BeaconAction."""
node_name: str
c2_server_ip_address: str
c2_server_ip_address: StrIP
keep_alive_frequency: int = Field(default=5, ge=1)
masquerade_protocol: str = Field(default="TCP")
masquerade_port: str = Field(default="HTTP")
masquerade_protocol: IPProtocol = Field(default="tcp")
masquerade_port: Port = Field(default=80)
@classmethod
def form_request(self, config: ConfigSchema) -> RequestFormat:
@@ -121,7 +124,7 @@ class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_c
"""Configuration schema for NodeSendRemoteCommandAction."""
node_name: str
remote_ip: str
remote_ip: StrIP
command: RequestFormat
@classmethod
@@ -149,7 +152,7 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
node_name: str
commands: Union[List[RequestFormat], RequestFormat]
ip_address: Optional[str]
ip_address: Optional[StrIP]
username: Optional[str]
password: Optional[str]
@@ -198,7 +201,7 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
node_name: str
username: Optional[str]
password: Optional[str]
target_ip_address: str
target_ip_address: StrIP
target_file_name: str
target_folder_name: str
exfiltration_folder_name: Optional[str]
@@ -229,7 +232,7 @@ class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_databa
"""Schema for options that can be passed to this action."""
node_name: str
server_ip_address: Optional[str] = None
server_ip_address: Optional[StrIP] = None
server_password: Optional[str] = None
@classmethod

22
src/primaite/game/agent/observations/acl_observation.py
View File

@@ -1,7 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Dict, List, Optional
from gymnasium import spaces
@@ -10,6 +9,9 @@ from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.ipv4_address import StrIP
from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
@@ -20,13 +22,13 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for ACLObservation."""
ip_list: Optional[List[IPv4Address]] = None
ip_list: Optional[List[StrIP]] = None
"""List of IP addresses."""
wildcard_list: Optional[List[str]] = None
"""List of wildcard strings."""
port_list: Optional[List[str]] = None
port_list: Optional[List[Port]] = None
"""List of port names."""
protocol_list: Optional[List[str]] = None
protocol_list: Optional[List[IPProtocol]] = None
"""List of protocol names."""
num_rules: Optional[int] = None
"""Number of ACL rules."""
@@ -35,10 +37,10 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
self,
where: WhereType,
num_rules: int,
ip_list: List[IPv4Address],
ip_list: List[StrIP],
wildcard_list: List[str],
port_list: List[str],
protocol_list: List[str],
port_list: List[Port],
protocol_list: List[IPProtocol],
) -> None:
"""
Initialise an ACL observation instance.
@@ -48,13 +50,13 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
:param num_rules: Number of ACL rules.
:type num_rules: int
:param ip_list: List of IP addresses.
:type ip_list: List[IPv4Address]
:type ip_list: List[StrIP]
:param wildcard_list: List of wildcard strings.
:type wildcard_list: List[str]
:param port_list: List of port names.
:type port_list: List[str]
:type port_list: List[Port]
:param protocol_list: List of protocol names.
:type protocol_list: List[str]
:type protocol_list: List[IPProtocol]
"""
self.where = where
self.num_rules: int = num_rules

21
src/primaite/game/agent/observations/firewall_observation.py
View File

@@ -11,6 +11,9 @@ from primaite.game.agent.observations.acl_observation import ACLObservation
from primaite.game.agent.observations.nic_observations import PortObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.ipv4_address import StrIP
from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
@@ -23,13 +26,13 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
hostname: str
"""Hostname of the firewall node, used for querying simulation state dictionary."""
ip_list: Optional[List[str]] = None
ip_list: Optional[List[StrIP]] = None
"""List of IP addresses for encoding ACLs."""
wildcard_list: Optional[List[str]] = None
"""List of IP wildcards for encoding ACLs."""
port_list: Optional[List[str]] = None
port_list: Optional[List[Port]] = None
"""List of ports for encoding ACLs."""
protocol_list: Optional[List[str]] = None
protocol_list: Optional[List[IPProtocol]] = None
"""List of protocols for encoding ACLs."""
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""
@@ -39,10 +42,10 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
def __init__(
self,
where: WhereType,
ip_list: List[str],
ip_list: List[StrIP],
wildcard_list: List[str],
port_list: List[str],
protocol_list: List[str],
port_list: List[Port],
protocol_list: List[IPProtocol],
num_rules: int,
include_users: bool,
) -> None:
@@ -53,13 +56,13 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
A typical location for a firewall might be ['network', 'nodes', <firewall_hostname>].
:type where: WhereType
:param ip_list: List of IP addresses.
:type ip_list: List[str]
:type ip_list: List[StrIP]
:param wildcard_list: List of wildcard rules.
:type wildcard_list: List[str]
:param port_list: List of port names.
:type port_list: List[str]
:type port_list: List[Port]
:param protocol_list: List of protocol types.
:type protocol_list: List[str]
:type protocol_list: List[IPProtocol]
:param num_rules: Number of rules configured in the firewall.
:type num_rules: int
:param include_users: If True, report user session information.

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

@@ -13,6 +13,7 @@ from primaite.game.agent.observations.host_observations import HostObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.observations.router_observation import RouterObservation
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.ipv4_address import StrIP
from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
@@ -52,13 +53,13 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
"""If True, report user session information."""
num_ports: Optional[int] = None
"""Number of ports."""
ip_list: Optional[List[str]] = None
ip_list: Optional[List[StrIP]] = None
"""List of IP addresses for encoding ACLs."""
wildcard_list: Optional[List[str]] = None
"""List of IP wildcards for encoding ACLs."""
port_list: Optional[List[str]] = None
port_list: Optional[List[Port]] = None
"""List of ports for encoding ACLs."""
protocol_list: Optional[List[str]] = None
protocol_list: Optional[List[IPProtocol]] = None
"""List of protocols for encoding ACLs."""
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""

9
src/primaite/game/agent/observations/router_observation.py
View File

@@ -11,6 +11,9 @@ from primaite.game.agent.observations.acl_observation import ACLObservation
from primaite.game.agent.observations.nic_observations import PortObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.ipv4_address import StrIP
from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
@@ -29,13 +32,13 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
"""Number of port observations configured for this router."""
acl: Optional[ACLObservation.ConfigSchema] = None
"""Configuration of ACL observation on this router."""
ip_list: Optional[List[str]] = None
ip_list: Optional[List[StrIP]] = None
"""List of IP addresses for encoding ACLs."""
wildcard_list: Optional[List[str]] = None
"""List of IP wildcards for encoding ACLs."""
port_list: Optional[List[str]] = None
port_list: Optional[List[Port]] = None
"""List of ports for encoding ACLs."""
protocol_list: Optional[List[str]] = None
protocol_list: Optional[List[IPProtocol]] = None
"""List of protocols for encoding ACLs."""
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""

6
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -83,6 +83,12 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
next_execution_timestep: int = 0
"""Timestep of the next action execution by the agent."""
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self._set_next_execution_timestep(
timestep=self.config.agent_settings.start_step, variance=self.config.agent_settings.start_variance
)
@computed_field
@cached_property
def start_node(self) -> str: