#2912 - End of day commit

This commit is contained in:
Charlie Crane
2024-10-17 16:59:44 +01:00
parent cd30e2d084
commit a90aec2bcd
9 changed files with 446 additions and 8 deletions

View File

@@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, field_validator, ValidationInfo
from primaite.game.agent.actions.manager import AbstractAction
from primaite.game.game import _LOGGER
from primaite.interface.request import RequestFormat
class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
@@ -168,3 +169,173 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
dst_port,
position,
]
class RouterACLRemoveRuleAction(AbstractAction):
"""Action which removes a rule from a router's ACL."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for RouterACLRemoveRuleAction."""
target_router: str
position: str
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
class FirewallACLAddRuleAction(AbstractAction):
"""Action which adds a rule to a firewall port's ACL."""
def __init__(
self,
manager: "ActionManager",
max_acl_rules: int,
num_ips: int,
num_ports: int,
num_protocols: int,
**kwargs,
) -> None:
"""Init method for FirewallACLAddRuleAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
:param num_ips: Number of IP addresses in the simulation.
:type num_ips: int
:param num_ports: Number of ports in the simulation.
:type num_ports: int
:param num_protocols: Number of protocols in the simulation.
:type num_protocols: int
"""
super().__init__(manager=manager)
num_permissions = 3
self.shape: Dict[str, int] = {
"position": max_acl_rules,
"permission": num_permissions,
"source_ip_id": num_ips,
"dest_ip_id": num_ips,
"source_port_id": num_ports,
"dest_port_id": num_ports,
"protocol_id": num_protocols,
}
def form_request(
self,
target_firewall_nodename: str,
firewall_port_name: str,
firewall_port_direction: str,
position: int,
permission: int,
source_ip_id: int,
source_wildcard_id: int,
dest_ip_id: int,
dest_wildcard_id: int,
source_port_id: int,
dest_port_id: int,
protocol_id: int,
) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if permission == 0:
permission_str = "UNUSED"
return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
elif permission == 1:
permission_str = "PERMIT"
elif permission == 2:
permission_str = "DENY"
else:
_LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.")
if protocol_id == 0:
return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
if protocol_id == 1:
protocol = "ALL"
else:
protocol = self.manager.get_internet_protocol_by_idx(protocol_id - 2)
# subtract 2 to account for UNUSED=0 and ALL=1.
if source_ip_id == 0:
return ["do_nothing"] # invalid formulation
elif source_ip_id == 1:
src_ip = "ALL"
else:
src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if source_port_id == 0:
return ["do_nothing"] # invalid formulation
elif source_port_id == 1:
src_port = "ALL"
else:
src_port = self.manager.get_port_by_idx(source_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if dest_ip_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_ip_id == 1:
dst_ip = "ALL"
else:
dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if dest_port_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_port_id == 1:
dst_port = "ALL"
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)
return [
"network",
"node",
target_firewall_nodename,
firewall_port_name,
firewall_port_direction,
"acl",
"add_rule",
permission_str,
protocol,
str(src_ip),
src_wildcard,
src_port,
str(dst_ip),
dst_wildcard,
dst_port,
position,
]
class FirewallACLRemoveRuleAction(AbstractAction):
"""Action which removes a rule from a firewall port's ACL."""
def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None:
"""Init method for RouterACLRemoveRuleAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"position": max_acl_rules}
def form_request(
self, target_firewall_nodename: str, firewall_port_name: str, firewall_port_direction: str, position: int
) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
target_firewall_nodename,
firewall_port_name,
firewall_port_direction,
"acl",
"remove_rule",
position,
]

View File

@@ -1,7 +1,5 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import ClassVar, Dict
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
@@ -16,6 +14,8 @@ class NodeApplicationAbstractAction(AbstractAction):
"""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for Node Application actions."""
node_name: str
application_name: str
@@ -33,6 +33,8 @@ class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="no
"""Action which executes an application."""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationExecuteAction."""
verb: str = "execute"
@@ -40,6 +42,8 @@ class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_
"""Action which scans an application."""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationScanAction."""
verb: str = "scan"
@@ -47,6 +51,8 @@ class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node
"""Action which closes an application."""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationCloseAction."""
verb: str = "close"
@@ -54,11 +60,28 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_a
"""Action which fixes an application."""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationFixAction."""
verb: str = "fix"
class NodeApplicationInstallAction(AbstractAction):
class NodeApplicationInstallAction(NodeApplicationAbstractAction):
"""Action which installs an application."""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationInstallAction."""
verb: str = "install"
# TODO: Either changes to application form_request bits, or add that here.
class NodeApplicationRemoveAction(NodeApplicationAbstractAction):
"""Action which removes/uninstalls an application"""
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationRemoveAction."""
verb: str = "uninstall"
# TODO: Either changes to application form_request bits, or add that here.

View File

@@ -0,0 +1,107 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Dict, Optional
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
class ConfigureRansomwareScriptAction(AbstractAction):
"""Action which sets config parameters for a ransomware script on a node."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this option."""
model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
payload: Optional[str] = None
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, config: Dict) -> 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:
return ["do_nothing"]
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "RansomwareScript", "configure", config]
class ConfigureDoSBotAction(AbstractAction):
"""Action which sets config parameters for a DoS bot on a node."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
model_config = ConfigDict(extra="forbid")
target_ip_address: Optional[str] = None
target_port: Optional[str] = None
payload: Optional[str] = None
repeat: Optional[bool] = None
port_scan_p_of_success: Optional[float] = None
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, node_id: int, config: Dict) -> 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:
return ["do_nothing"]
self._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "DoSBot", "configure", config]
class ConfigureC2BeaconAction(AbstractAction):
"""Action which configures a C2 Beacon based on the parameters given."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for ConfigureC2BeaconAction."""
node_name: str
c2_server_ip_address: str
keep_alive_frequency: int = Field(default=5, ge=1)
masquerade_protocol: str = Field(default="TCP")
masquerade_port: str = Field(default="HTTP")
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
c2_server_ip_address: str
keep_alive_frequency: int = Field(default=5, ge=1)
masquerade_protocol: str = Field(default="TCP")
masquerade_port: str = Field(default="HTTP")
@field_validator(
"c2_server_ip_address",
"keep_alive_frequency",
"masquerade_protocol",
"masquerade_port",
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(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"]
config = ConfigureC2BeaconAction._Opts(
c2_server_ip_address=config["c2_server_ip_address"],
keep_alive_frequency=config["keep_alive_frequency"],
masquerade_port=config["masquerade_port"],
masquerade_protocol=config["masquerade_protocol"],
)
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config.__dict__]

View File

@@ -8,11 +8,13 @@ from primaite.interface.request import RequestFormat
class NodeFileAbstractAction(AbstractAction):
"""Abstract base class for file actions.
Any action which applies to a file and uses node_name, folder_name, and file_name as its only three parameters can inherit
from this base class.
Any action which applies to a file and uses node_name, folder_name, and file_name as its
only three parameters can inherit from this base class.
"""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for NodeFileAbstractAction."""
node_name: str
folder_name: str
file_name: str
@@ -41,6 +43,8 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
"""Action which creates a new file in a given folder."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileCreateAction."""
verb: str = "create"
@@ -48,6 +52,8 @@ class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
"""Action which scans a file."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileScanAction."""
verb: str = "scan"
@@ -55,6 +61,8 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
"""Action which deletes a file."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileDeleteAction."""
verb: str = "delete"
@@ -62,6 +70,8 @@ class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restor
"""Action which restores a file."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileRestoreAction."""
verb: str = "restore"
@@ -69,6 +79,8 @@ class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrup
"""Action which corrupts a file."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileCorruptAction."""
verb: str = "corrupt"
@@ -76,4 +88,6 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
"""Action which increases a file's access count."""
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileAccessAction."""
verb: str = "access"

View File

@@ -15,6 +15,8 @@ class NodeFolderAbstractAction(AbstractAction):
"""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base configuration schema for NodeFolder actions."""
node_name: str
folder_name: str
@@ -34,6 +36,8 @@ class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_sca
"""Action which scans a folder."""
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderScanAction."""
verb: str = "scan"
@@ -41,6 +45,8 @@ class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folde
"""Action which checks the hash of a folder."""
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderCheckhashAction."""
verb: str = "checkhash"
@@ -48,6 +54,8 @@ class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_r
"""Action which repairs a folder."""
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderRepairAction."""
verb: str = "repair"
@@ -55,6 +63,8 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_
"""Action which restores a folder."""
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderRestoreAction."""
verb: str = "restore"
@@ -62,4 +72,6 @@ class NodeFolderCreateAction(AbstractAction, identifier="node_folder_create"):
"""Action which creates a new folder."""
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderCreateAction."""
verb: str = "create"

View File

@@ -0,0 +1,57 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Dict, Optional
from pydantic import BaseModel, ConfigDict
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
class HostNICAbstractAction(AbstractAction):
"""
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.
"""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for HostNIC actions."""
num_nodes: str
max_nics_per_node: str
node_name: str
nic_num: str
def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
"""Init method for HostNICAbstractAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param num_nodes: Number of nodes in the simulation.
:type num_nodes: int
:param max_nics_per_node: Maximum number of NICs per node.
:type max_nics_per_node: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
self.verb: str # define but don't initialise: defends against children classes not defining this
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.nic_num is None:
return ["do_nothing"]
return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb]
class HostNICEnableAction(HostNICAbstractAction):
"""Action which enables a NIC."""
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
"""Configuration schema for HostNICEnableAction."""
verb: str = "enable"
class HostNICDisableAction(HostNICAbstractAction):
"""Action which disables a NIC."""
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
"""Configuration schema for HostNICDisableAction."""
verb: str = "disable"

View File

@@ -14,13 +14,13 @@ agents:
from __future__ import annotations
import itertools
from abc import ABC, abstractmethod
from abc import ABC
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type
from gymnasium import spaces
from pydantic import BaseModel, ConfigDict
from primaite.game.game import _LOGGER, PrimaiteGame
from primaite.game.game import PrimaiteGame
from primaite.interface.request import RequestFormat
@@ -57,6 +57,7 @@ class DoNothingAction(AbstractAction):
type: Literal["do_nothing"] = "do_nothing"
def form_request(self, options: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["do_nothing"]

View File

@@ -0,0 +1,43 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Dict, Optional
from pydantic import BaseModel, ConfigDict
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
class NetworkPortEnableAction(AbstractAction):
"""Action which enables are port on a router or a firewall."""
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
"""Init method for NetworkPortEnableAction.
:param max_nics_per_node: Maximum number of NICs per node.
:type max_nics_per_node: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"port_id": max_nics_per_node}
def form_request(self, target_nodename: str, port_id: int) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if target_nodename is None or port_id is None:
return ["do_nothing"]
return ["network", "node", target_nodename, "network_interface", port_id, "enable"]
class NetworkPortDisableAction(AbstractAction):
"""Action which disables are port on a router or a firewall."""
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
"""Init method for NetworkPortDisableAction.
:param max_nics_per_node: Maximum number of NICs per node.
:type max_nics_per_node: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"port_id": max_nics_per_node}
def form_request(self, target_nodename: str, port_id: int) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if target_nodename is None or port_id is None:
return ["do_nothing"]
return ["network", "node", target_nodename, "network_interface", port_id, "disable"]

View File

@@ -14,6 +14,8 @@ class NodeAbstractAction(AbstractAction):
"""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for Node actions."""
node_name: str
verb: ClassVar[str]
@@ -28,6 +30,8 @@ class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
"""Action which scans a node's OS."""
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeOSScanAction."""
verb: str = "scan"
@@ -35,6 +39,8 @@ class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
"""Action which shuts down a node."""
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeShutdownAction."""
verb: str = "shutdown"
@@ -42,6 +48,8 @@ class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
"""Action which starts up a node."""
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeStartupAction."""
verb: str = "startup"
@@ -49,4 +57,6 @@ class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
"""Action which resets a node."""
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeResetAction."""
verb: str = "reset"