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

#2912 - Updated to include __all__ and some test fixes. New extensible actions documentation page

Browse Source
This commit is contained in:
Charlie Crane
2024-11-12 14:49:44 +00:00
parent 844a3a60fa
commit d757bd01f0
15 changed files with 239 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -175,6 +175,10 @@ class NodeApplicationInstallAction(AbstractAction):
class ConfigureDatabaseClientAction(AbstractAction):
"""Action which sets config parameters for a database client on a node."""
model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
@@ -197,9 +201,10 @@ class ConfigureDatabaseClientAction(AbstractAction):
class ConfigureRansomwareScriptAction(AbstractAction):
"""Action which sets config parameters for a ransomware script on a node."""
class _Opts(BaseModel):
class _Opts(BaseModel, AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this option."""
node_name: str
model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
@@ -208,13 +213,20 @@ class ConfigureRansomwareScriptAction(AbstractAction):
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, config: Dict) -> RequestFormat:
def form_request(self, config: _Opts) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
if node_name is None:
if config.node_name is None:
return ["do_nothing"]
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "RansomwareScript", "configure", config]
return [
"network",
"node",
config.node_name,
"application",
"RansomwareScript",
"configure",
config.model_config,
]
class ConfigureDoSBotAction(AbstractAction):
@@ -285,8 +297,8 @@ class NodeFolderAbstractAction(AbstractAction):
class NodeFolderScanAction(NodeFolderAbstractAction):
"""Action which scans a folder."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
def __init__(self, manager: "ActionManager", node_name: str, folder_name, **kwargs) -> None:
super().__init__(manager, node_name=node_name, folder_name=folder_name, **kwargs)
self.verb: str = "scan"

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

@@ -17,8 +17,8 @@ from primaite.game.agent.actions import (
from primaite.game.agent.actions.manager import ActionManager
__all__ = (
"acl",
"abstract",
"acl",
"application",
"config",
"file",

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

@@ -8,17 +8,18 @@ from pydantic import BaseModel, ConfigDict
from primaite.interface.request import RequestFormat
# notes:
# we actually don't need to hold any state in actions, so there's no need to define any __init__ logic.
# all the init methods in the old actions are just used for holding a verb and shape, which are not really used.
# the config schema should be used to the actual parameters for formatting the action itself.
# (therefore there's no need for creating action instances, just the action class contains logic for converting
# CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be
# instantiated.)
class AbstractAction(BaseModel):
"""Base class for actions."""
# notes:
# we actually don't need to hold any state in actions, so there's no need to define any __init__ logic.
# all the init methods in the old actions are just used for holding a verb and shape, which are not really used.
# the config schema should be used to the actual parameters for formatting the action itself.
# (therefore there's no need for creating action instances, just the action class contains logic for converting
# CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be
# instantiated.)
class ConfigSchema(BaseModel, ABC):
"""Base configuration schema for Actions."""
@@ -41,12 +42,7 @@ class AbstractAction(BaseModel):
@classmethod
def from_config(cls, config: Dict) -> "AbstractAction":
"""Create an action component from a config dictionary."""
# set attributes for action based off config dict
# if config["type"] not in cls._registry:
# raise ValueError(f"Invalid action reward type {config['type']}")
for attribute, value in config.items():
if not hasattr(cls.ConfigSchema, attribute):
setattr(cls.ConfigSchema, attribute, value)
return cls

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

@@ -17,27 +17,35 @@ __all__ = (
class ACLAbstractAction(AbstractAction, identifier="acl_abstract_action"):
"""Base class for ACL actions."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema base for ACL abstract actions."""
class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
"""Action which adds a rule to a router's ACL."""
target_router: str
position: int
permission: Literal[1, 2]
src_ip: str
source_wildcard_id: int
source_port: str
dst_ip: str
dst_wildcard: int
dst_port: int
protocol_name: str
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for RouterACLAddRuleAction."""
target_router: str
position: int
permission: Literal[1, 2]
source_ip_id: int
source_wildcard_id: int
source_port_id: int
src_ip: str
src_wildcard: int
source_port: str
dst_ip: str
dst_wildcard_id: int
dst_wildcard: int
dst_port: int
protocol_name: str
@@ -50,95 +58,54 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
"""At what position to add the rule, must be specified."""
permission: Literal[1, 2]
"""Whether to allow or deny traffic, must be specified. 1 = PERMIT, 2 = DENY."""
source_ip_id: int = Field(default=1, ge=1)
src_ip: str
"""Rule source IP address. By default, all ip addresses."""
source_wildcard_id: int = Field(default=0, ge=0)
"""Rule source IP wildcard. By default, use the wildcard at index 0 from action manager."""
source_port_id: int = Field(default=1, ge=1)
source_port: int = Field(default=1, ge=1)
"""Rule source port. By default, all source ports."""
dst_ip_id: int = Field(default=1, ge=1)
"""Rule destination IP address. By default, all ip addresses."""
dst_wildcard_id: int = Field(default=0, ge=0)
dst_wildcard: int = Field(default=0, ge=0)
"""Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager."""
dst_port_id: int = Field(default=1, ge=1)
"""Rule destination port. By default, all destination ports."""
protocol_name: str = "ALL"
"""Rule protocol. By default, all protocols."""
@field_validator(
"source_ip_id",
"source_port_id",
"source_wildcard_id",
"dest_ip_id",
"dest_port_id",
"dest_wildcard_id",
"protocol_name",
mode="before",
)
@classmethod
def not_none(cls, v: str, info: ValidationInfo) -> int:
"""If None is passed, use the default value instead."""
if v is None:
return cls.model_fields[info.field_name].default
return v
# @field_validator(
# "source_ip_id",
# "source_port_id",
# "source_wildcard_id",
# "dest_ip_id",
# "dest_port_id",
# "dest_wildcard_id",
# "protocol_name",
# mode="before",
# )
# @classmethod
# def not_none(cls, v: str, info: ValidationInfo) -> int:
# """If None is passed, use the default value instead."""
# if v is None:
# return cls.model_fields[info.field_name].default
# return v
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
# Validate incoming data.
parsed_options = RouterACLAddRuleAction.ACLRuleOptions(
target_router=config.target_router,
position=config.position,
permission=config.permission,
source_ip_id=config.source_ip_id,
source_wildcard_id=config.source_wildcard_id,
dest_ip_id=config.dst_ip,
dest_wildcard_id=config.dest_wildcard_id,
source_port_id=config.source_port_id,
dest_port_id=config.dst_port_id,
protocol=config.protocol_name,
)
if parsed_options.permission == 1:
permission_str = "PERMIT"
elif parsed_options.permission == 2:
permission_str = "DENY"
# else:
# _LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.")
if parsed_options.protocol_name == "ALL":
protocol = "ALL"
else:
protocol = parsed_options.protocol_name
# subtract 2 to account for UNUSED=0 and ALL=1.
if parsed_options.source_ip_id == 1:
src_ip = "ALL"
else:
src_ip = cls.manager.get_ip_address_by_idx(parsed_options.source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
src_wildcard = cls.manager.get_wildcard_by_idx(parsed_options.source_wildcard_id)
if parsed_options.source_port_id == 1:
src_port = "ALL"
else:
src_port = cls.manager.get_port_by_idx(parsed_options.source_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if parsed_options.dest_ip_id == 1:
dst_ip = "ALL"
else:
dst_ip = cls.manager.get_ip_address_by_idx(parsed_options.dest_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
dst_ip=config.dst_ip
dst_wildcard = config.dest_wildcard_id
if parsed_options.dest_port_id == 1:
dst_port = "ALL"
else:
dst_port = cls.manager.get_port_by_idx(parsed_options.dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
# parsed_options = RouterACLAddRuleAction.ACLRuleOptions(
# target_router=config.target_router,
# position=config.position,
# permission=config.permission,
# src_ip=config.src_ip,
# source_wildcard_id=config.source_wildcard_id,
# dst_ip_id=config.dst_ip,
# dst_wildcard=config.dst_wildcard,
# source_port_id=config.source_port_id,
# dest_port=config.dst_port,
# protocol=config.protocol_name,
# )
return [
"network",
@@ -146,11 +113,11 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
config.target_router,
"acl",
"add_rule",
config.permission_str,
protocol,
str(src_ip),
config.permission,
config.protocol_name,
config.src_ip,
config.src_wildcard,
config.src_port,
config.source_port,
str(config.dst_ip),
config.dst_wildcard,
config.dst_port,
@@ -193,7 +160,6 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
num_permissions: int = 3
permission: str
def __init__(
self,
manager: "ActionManager",
@@ -228,10 +194,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
"protocol_id": num_protocols,
}
@classmethod
def form_request(cls, config:ConfigSchema) -> List[str]:
def form_request(cls, config: ConfigSchema) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.permission == 0:
permission_str = "UNUSED"
@@ -260,6 +224,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
else:
# src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
pass
if config.source_port_id == 0:
return ["do_nothing"] # invalid formulation
@@ -286,8 +251,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
else:
# dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
# src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
# dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
# src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
# dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
pass
return [
@@ -298,7 +263,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
config.firewall_port_direction,
"acl",
"add_rule",
permission_str,
config.permission,
protocol,
str(src_ip),
config.src_wildcard,

24
src/primaite/game/agent/actions/config.py
View File

@@ -15,6 +15,7 @@ __all__ = (
"TerminalC2ServerAction",
"RansomwareLaunchC2ServerAction",
"ExfiltrationC2ServerAction",
"ConfigureDatabaseClientAction",
)
@@ -230,3 +231,26 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
}
ExfiltrationC2ServerAction._Opts.model_validate(command_model)
return ["network", "node", node_name, "application", "C2Server", "exfiltrate", command_model]
class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"):
"""Action which sets config parameters for a database client on a node."""
node_name: str
model_config: ConfigDict = ConfigDict(extra="forbid")
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
model_config = ConfigDict(extra="forbid")
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, 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]

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

@@ -11,6 +11,8 @@ __all__ = (
"NodeFileRestoreAction",
"NodeFileCorruptAction",
"NodeFileAccessAction",
"NodeFileCheckhashAction",
"NodeFileRepairAction",
)

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

@@ -24,6 +24,8 @@ from primaite.interface.request import RequestFormat
# TODO: Make sure that actions are backwards compatible where the old YAML format is used.
__all__ = "DoNothingAction"
class DoNothingAction(AbstractAction, identifier="do_nothing"):
"""Do Nothing Action."""
@@ -69,7 +71,7 @@ class ActionManager:
Action mapping that converts an integer to a specific action and parameter choice.
For example :
{0: ("NODE_SERVICE_SCAN", {node_id:0, service_id:2})}
{0: ("node_service_scan", {node_name:"client_1", service_name:"WebBrowser"})}
"""
if act_map is None:
# raise RuntimeError("Action map must be specified in the config file.")
@@ -147,7 +149,7 @@ class ActionManager:
Since the agent uses a discrete action space which acts as a flattened version of the component-based
action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful
action and values of parameters. For example action 0 can correspond to do nothing, action 1 can
correspond to "NODE_SERVICE_SCAN" with ``node_id=1`` and ``service_id=1``, action 2 can be "
correspond to "node_service_scan" with ``node_name="server"`` and ``service_name="WebBrowser"``, action 2 can be "
3. ``options``
``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method.
These options are used to calculate the shape of the action space, and to provide additional information

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

@@ -5,7 +5,15 @@ from typing import ClassVar, List, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = ("NodeOSScanAction", "NodeShutdownAction", "NodeStartupAction", "NodeResetAction")
__all__ = (
"NodeOSScanAction",
"NodeShutdownAction",
"NodeStartupAction",
"NodeResetAction",
"NodeNMAPPingScanAction",
"NodeNMAPPortScanAction",
"NodeNetworkServiceReconAction",
)
class NodeAbstractAction(AbstractAction, identifier="node_abstract"):

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

@@ -4,7 +4,11 @@ from abc import abstractmethod
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = ("NodeSessionsRemoteLoginAction", "NodeSessionsRemoteLogoutAction")
__all__ = (
"NodeSessionsRemoteLoginAction",
"NodeSessionsRemoteLogoutAction",
"NodeAccountChangePasswordAction",
)
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"):
@@ -73,7 +77,7 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node
return ["network", "node", config.node_name, "service", "Terminal", "remote_logoff", config.remote_ip]
class NodeAccountsChangePasswordAction(NodeSessionAbstractAction, identifier="node_accounts_change_password"):
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"):
"""Action which changes the password for a user."""
username: str