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

update more tests to new actions schema

Browse Source
This commit is contained in:
Marek Wolan
2025-01-17 14:38:59 +00:00
parent 03cab0fcec
commit 858406c4a3
48 changed files with 351 additions and 413 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/game/agent/actions/abstract.py
View File

@@ -18,7 +18,7 @@ class AbstractAction(BaseModel, ABC):
"""Base configuration schema for Actions."""
model_config = ConfigDict(extra="forbid")
type: str
type: str = ""
_registry: ClassVar[Dict[str, Type[AbstractAction]]] = {}

10
src/primaite/game/agent/actions/acl.py
View File

@@ -28,7 +28,7 @@ class ACLAddRuleAbstractAction(AbstractAction, ABC):
src_ip: IPV4Address
protocol_name: Union[IPProtocol, Literal["ALL"]]
permission: Literal["ALLOW", "DENY"]
permission: Literal["PERMIT", "DENY"]
position: int
dst_ip: Union[IPV4Address, Literal["ALL"]]
src_port: Union[Port, Literal["ALL"]]
@@ -70,10 +70,10 @@ class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_ad
config.permission,
config.protocol_name,
str(config.src_ip),
config.src_wildcard,
str(config.src_wildcard),
config.src_port,
str(config.dst_ip),
config.dst_wildcard,
str(config.dst_wildcard),
config.dst_port,
config.position,
]
@@ -121,10 +121,10 @@ class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_ac
config.permission,
config.protocol_name,
str(config.src_ip),
config.src_wildcard,
str(config.src_wildcard),
config.src_port,
str(config.dst_ip),
config.dst_wildcard,
str(config.dst_wildcard),
config.dst_port,
config.position,
]

2
src/primaite/game/agent/actions/application.py
View File

@@ -19,7 +19,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
"""
Base class for application actions.
Any action which applies to an application and uses node_id and application_id as its only two parameters can
Any action which applies to an application and uses node_name and application_name as its only two parameters can
inherit from this base class.
"""

4
src/primaite/game/agent/actions/host_nic.py
View File

@@ -12,8 +12,8 @@ class HostNICAbstractAction(AbstractAction, ABC):
"""
Abstract base class for NIC actions.
Any action which applies to a NIC and uses node_id and nic_id as its only two parameters can inherit from this base
class.
Any action which applies to a NIC and uses node_name and nic_num as its only two parameters can inherit from this
base class.
"""
config: "HostNICAbstractAction.ConfigSchema"

2
src/primaite/game/agent/actions/manager.py
View File

@@ -84,7 +84,7 @@ class ActionManager(BaseModel):
def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat:
"""Take action in CAOS format and use the execution definition to change it into PrimAITE request format."""
act_class = AbstractAction._registry[action_identifier]
config = act_class.ConfigSchema(type=action_identifier, **action_options)
config = act_class.ConfigSchema(**action_options)
return act_class.form_request(config=config)
@property

23
src/primaite/game/agent/actions/node.py
View File

@@ -4,6 +4,8 @@ from typing import ClassVar, List, Optional, Union
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.port import Port
__all__ = (
"NodeOSScanAction",
@@ -92,7 +94,7 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti
target_ip_address: Union[str, List[str]]
show: bool = False
node_name: str
source_node: str
@classmethod
@abstractmethod
@@ -107,18 +109,13 @@ class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_
config: "NodeNMAPPingScanAction.ConfigSchema"
class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema):
"""Configuration schema for NodeNMAPPingScanAction."""
pass
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]: # noqa
def form_request(cls, config: "NodeNMAPPingScanAction.ConfigSchema") -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.node_name,
config.source_node,
"application",
"NMAP",
"ping_scan",
@@ -135,8 +132,8 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
"""Configuration Schema for NodeNMAPPortScanAction."""
source_node: str
target_protocol: Optional[Union[str, List[str]]] = (None,)
target_port: Optional[Union[str, List[str]]] = (None,)
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None
target_port: Optional[Union[Port, List[Port]]] = None
show: Optional[bool] = (False,)
@classmethod
@@ -166,11 +163,11 @@ class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_net
config: "NodeNetworkServiceReconAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema):
"""Configuration schema for NodeNetworkServiceReconAction."""
target_protocol: Optional[Union[str, List[str]]] = (None,)
target_port: Optional[Union[str, List[str]]] = (None,)
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None
target_port: Optional[Union[Port, List[Port]]] = None
show: Optional[bool] = (False,)
@classmethod

2
src/primaite/game/agent/actions/session.py
View File

@@ -55,7 +55,7 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_
config.node_name,
"service",
"Terminal",
"ssh_to_remote",
"node_session_remote_login",
config.username,
config.password,
config.remote_ip,

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

@@ -4,7 +4,7 @@ from typing import List, Optional, Union
from pydantic import ConfigDict, Field, field_validator, ValidationInfo
from primaite.game.agent.actions.manager import AbstractAction, ActionManager
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
@@ -28,36 +28,31 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="c2_server_rans
"""Configuration schema for ConfigureRansomwareScriptAction."""
node_name: str
server_ip_address: Optional[str]
server_password: Optional[str]
payload: Optional[str]
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
payload: Optional[str] = None
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"application",
"RansomwareScript",
"configure",
config.model_config,
]
data = dict(
server_ip_address=config.server_ip_address,
server_password=config.server_password,
payload=config.payload,
)
return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data]
class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
"""Action which sets config parameters for a DoS bot on a node."""
config: "ConfigureDoSBotAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
model_config = ConfigDict(extra="forbid")
node_name: str
target_ip_address: Optional[str] = None
target_port: Optional[str] = None
payload: Optional[str] = None
@@ -66,22 +61,24 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
dos_intensity: Optional[float] = None
max_sessions: Optional[int] = None
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, config: ConfigSchema) -> RequestFormat:
@classmethod
def form_request(config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
self.ConfigSchema.model_validate(config) # check that options adhere to schema
return ["network", "node", config.node_name, "application", "DoSBot", "configure", config]
data = dict(
target_ip_address=config.target_ip_address,
target_port=config.target_port,
payload=config.payload,
repeat=config.repeat,
port_scan_p_of_success=config.port_scan_p_of_success,
dos_intensity=config.dos_intensity,
max_sessions=config.max_sessions,
)
return ["network", "node", config.node_name, "application", "DoSBot", "configure", data]
class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
"""Action which configures a C2 Beacon based on the parameters given."""
config: "ConfigureC2BeaconAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for ConfigureC2BeaconAction."""
@@ -91,6 +88,7 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
masquerade_protocol: str = Field(default="TCP")
masquerade_port: str = Field(default="HTTP")
# TODO: this validator should not be needed anymore, test what happens if removed.
@field_validator(
"c2_server_ip_address",
"keep_alive_frequency",
@@ -108,7 +106,13 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
@classmethod
def form_request(self, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config]
data = dict(
c2_server_ip_address=config.c2_server_ip_address,
keep_alive_frequency=config.keep_alive_frequency,
masquerade_protocol=config.masquerade_protocol,
masquerade_port=config.masquerade_port,
)
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data]
class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"):
@@ -228,11 +232,13 @@ class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_databa
"""Schema for options that can be passed to this action."""
node_name: str
model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", config.model_config]
data = {"server_ip_address": config.server_ip_address, "server_password": config.server_password}
return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", data]

4
src/primaite/game/agent/interface.py
View File

@@ -58,9 +58,9 @@ class AbstractAgent(BaseModel, ABC):
model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
type: str
ref: str
ref: str = ""
"""name of the agent."""
team: Optional[Literal["BLUE", "GREEN", "RED"]]
team: Optional[Literal["BLUE", "GREEN", "RED"]] = None
agent_settings: AbstractAgent.AgentSettingsSchema = Field(default=lambda: AbstractAgent.AgentSettingsSchema())
action_space: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema())
observation_space: ObservationManager.ConfigSchema = Field(

4
src/primaite/game/agent/rewards.py
View File

@@ -51,7 +51,7 @@ class AbstractReward(BaseModel):
class ConfigSchema(BaseModel, ABC):
"""Config schema for AbstractReward."""
type: str
type: str = ""
_registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {}
@@ -404,7 +404,7 @@ class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"):
:rtype: float
"""
if last_action_response.action == "do_nothing":
return self.do_nothing_penalty
return self.config.do_nothing_penalty
else:
return self.config.action_penalty

3
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -16,7 +16,6 @@ __all__ = "ProbabilisticAgent"
class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema())
rng: Generator = np.random.default_rng(np.random.randint(0, 65535))
class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema):
@@ -52,6 +51,8 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema()
)
config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema())
@property
def probabilities(self) -> Dict[str, int]:
"""Convenience method to view the probabilities of the Agent."""