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

#2682 Backport changes to core PrimAITE

Browse Source
This commit is contained in:
Marek Wolan
2024-06-25 11:04:52 +01:00
parent 4a81dc3b2c
commit 28dabad66b
52 changed files with 456 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/primaite/__init__.py
View File

@@ -203,3 +203,6 @@ def getLogger(name: str) -> Logger: # noqa
logger.setLevel(PRIMAITE_CONFIG["logging"]["log_level"])
return logger
DEFAULT_BANDWIDTH: Final[int] = 100

32
src/primaite/config/_package_data/data_manipulation.yaml
View File

@@ -476,7 +476,7 @@ agents:
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
permission: 2
source_ip_id: 7 # client 1
@@ -489,7 +489,7 @@ agents:
47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
permission: 2
source_ip_id: 8 # client 2
@@ -502,7 +502,7 @@ agents:
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
permission: 2
source_ip_id: 7 # client 1
@@ -515,7 +515,7 @@ agents:
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
permission: 2
source_ip_id: 8 # client 2
@@ -528,7 +528,7 @@ agents:
50: # old action num: 26
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
permission: 2
source_ip_id: 7 # client 1
@@ -541,7 +541,7 @@ agents:
51: # old action num: 27
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
permission: 2
source_ip_id: 8 # client 2
@@ -554,52 +554,52 @@ agents:
52: # old action num: 28
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 0
53: # old action num: 29
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
54: # old action num: 30
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
55: # old action num: 31
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
56: # old action num: 32
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
57: # old action num: 33
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
58: # old action num: 34
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
59: # old action num: 35
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 7
60: # old action num: 36
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 8
61: # old action num: 37
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 9
62: # old action num: 38
action: "HOST_NIC_DISABLE"

68
src/primaite/config/_package_data/data_manipulation_marl.yaml
View File

@@ -470,7 +470,7 @@ agents:
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
permission: 2
source_ip_id: 7 # client 1
@@ -483,7 +483,7 @@ agents:
47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
permission: 2
source_ip_id: 8 # client 2
@@ -496,7 +496,7 @@ agents:
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
permission: 2
source_ip_id: 7 # client 1
@@ -509,7 +509,7 @@ agents:
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
permission: 2
source_ip_id: 8 # client 2
@@ -522,7 +522,7 @@ agents:
50: # old action num: 26
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
permission: 2
source_ip_id: 7 # client 1
@@ -535,7 +535,7 @@ agents:
51: # old action num: 27
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
permission: 2
source_ip_id: 8 # client 2
@@ -548,52 +548,52 @@ agents:
52: # old action num: 28
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 0
53: # old action num: 29
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
54: # old action num: 30
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
55: # old action num: 31
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
56: # old action num: 32
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
57: # old action num: 33
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
58: # old action num: 34
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
59: # old action num: 35
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 7
60: # old action num: 36
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 8
61: # old action num: 37
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 9
62: # old action num: 38
action: "HOST_NIC_DISABLE"
@@ -834,10 +834,10 @@ agents:
- type: NODE_RESET
- type: ROUTER_ACL_ADDRULE
options:
target_router_nodename: router_1
target_router: router_1
- type: ROUTER_ACL_REMOVERULE
options:
target_router_nodename: router_1
target_router: router_1
- type: HOST_NIC_ENABLE
- type: HOST_NIC_DISABLE
@@ -1053,7 +1053,7 @@ agents:
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
permission: 2
source_ip_id: 7 # client 1
@@ -1066,7 +1066,7 @@ agents:
47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
permission: 2
source_ip_id: 8 # client 2
@@ -1079,7 +1079,7 @@ agents:
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
permission: 2
source_ip_id: 7 # client 1
@@ -1092,7 +1092,7 @@ agents:
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
permission: 2
source_ip_id: 8 # client 2
@@ -1105,7 +1105,7 @@ agents:
50: # old action num: 26
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
permission: 2
source_ip_id: 7 # client 1
@@ -1118,7 +1118,7 @@ agents:
51: # old action num: 27
action: "ROUTER_ACL_ADDRULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
permission: 2
source_ip_id: 8 # client 2
@@ -1131,52 +1131,52 @@ agents:
52: # old action num: 28
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 0
53: # old action num: 29
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 1
54: # old action num: 30
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 2
55: # old action num: 31
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 3
56: # old action num: 32
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 4
57: # old action num: 33
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 5
58: # old action num: 34
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 6
59: # old action num: 35
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 7
60: # old action num: 36
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 8
61: # old action num: 37
action: "ROUTER_ACL_REMOVERULE"
options:
target_router_nodename: router_1
target_router: router_1
position: 9
62: # old action num: 38
action: "HOST_NIC_DISABLE"

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

@@ -11,9 +11,10 @@ AbstractAction. The ActionManager is responsible for:
"""
import itertools
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
from typing import Dict, List, Literal, Optional, Tuple, TYPE_CHECKING, Union
from gymnasium import spaces
from pydantic import BaseModel, Field, field_validator, ValidationInfo
from primaite import getLogger
@@ -321,12 +322,12 @@ class NodeFileCreateAction(AbstractAction):
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
self.verb: str = "create"
def form_request(self, node_id: int, folder_name: str, file_name: str) -> List[str]:
def form_request(self, node_id: int, folder_name: str, file_name: str, force: Optional[bool] = False) -> List[str]:
"""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)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name]
return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name, force]
class NodeFolderCreateAction(AbstractAction):
@@ -493,6 +494,47 @@ class NodeResetAction(NodeAbstractAction):
class RouterACLAddRuleAction(AbstractAction):
"""Action which adds a rule to a router's ACL."""
class ACLRuleOptions(BaseModel):
"""Validator for ACL_ADD_RULE options."""
target_router: str
"""On which router to add the rule, must be specified."""
position: int
"""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)
"""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)
"""Rule source port. By default, all source ports."""
dest_ip_id: int = Field(default=1, ge=1)
"""Rule destination IP address. By default, all ip addresses."""
dest_wildcard_id: int = Field(default=0, ge=0)
"""Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager."""
dest_port_id: int = Field(default=1, ge=1)
"""Rule destination port. By default, all destination ports."""
protocol_id: int = Field(default=1, ge=1)
"""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_id",
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
def __init__(
self,
manager: "ActionManager",
@@ -529,7 +571,7 @@ class RouterACLAddRuleAction(AbstractAction):
def form_request(
self,
target_router_nodename: str,
target_router: str,
position: int,
permission: int,
source_ip_id: int,
@@ -541,62 +583,63 @@ class RouterACLAddRuleAction(AbstractAction):
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:
# Validate incoming data.
parsed_options = RouterACLAddRuleAction.ACLRuleOptions(
target_router=target_router,
position=position,
permission=permission,
source_ip_id=source_ip_id,
source_wildcard_id=source_wildcard_id,
dest_ip_id=dest_ip_id,
dest_wildcard_id=dest_wildcard_id,
source_port_id=source_port_id,
dest_port_id=dest_port_id,
protocol_id=protocol_id,
)
if parsed_options.permission == 1:
permission_str = "PERMIT"
elif permission == 2:
elif parsed_options.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:
if parsed_options.protocol_id == 1:
protocol = "ALL"
else:
protocol = self.manager.get_internet_protocol_by_idx(protocol_id - 2)
protocol = self.manager.get_internet_protocol_by_idx(parsed_options.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:
if parsed_options.source_ip_id == 1:
src_ip = "ALL"
else:
src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
src_ip = self.manager.get_ip_address_by_idx(parsed_options.source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
if source_port_id == 0:
return ["do_nothing"] # invalid formulation
elif source_port_id == 1:
src_wildcard = self.manager.get_wildcard_by_idx(parsed_options.source_wildcard_id)
if parsed_options.source_port_id == 1:
src_port = "ALL"
else:
src_port = self.manager.get_port_by_idx(source_port_id - 2)
src_port = self.manager.get_port_by_idx(parsed_options.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:
if parsed_options.dest_ip_id == 1:
dst_ip = "ALL"
else:
dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2)
dst_ip = self.manager.get_ip_address_by_idx(parsed_options.dest_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
dst_wildcard = self.manager.get_wildcard_by_idx(parsed_options.dest_wildcard_id)
if dest_port_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_port_id == 1:
if parsed_options.dest_port_id == 1:
dst_port = "ALL"
else:
dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
dst_port = self.manager.get_port_by_idx(parsed_options.dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
return [
"network",
"node",
target_router_nodename,
target_router,
"acl",
"add_rule",
permission_str,
@@ -625,9 +668,9 @@ class RouterACLRemoveRuleAction(AbstractAction):
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"position": max_acl_rules}
def form_request(self, target_router_nodename: str, position: int) -> List[str]:
def form_request(self, target_router: str, position: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", target_router_nodename, "acl", "remove_rule", position]
return ["network", "node", target_router, "acl", "remove_rule", position]
class FirewallACLAddRuleAction(AbstractAction):
@@ -877,7 +920,10 @@ class NodeNMAPPingScanAction(AbstractAction):
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, source_node: str, target_ip_address: Union[str, List[str]]) -> List[str]: # noqa
def form_request(
self, source_node: str, target_ip_address: Union[str, List[str]], show: Optional[bool] = False
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
@@ -885,7 +931,7 @@ class NodeNMAPPingScanAction(AbstractAction):
"application",
"NMAP",
"ping_scan",
{"target_ip_address": target_ip_address},
{"target_ip_address": target_ip_address, "show": show},
]
@@ -901,6 +947,7 @@ class NodeNMAPPortScanAction(AbstractAction):
target_ip_address: Union[str, List[str]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[str, List[str]]] = None,
show: Optional[bool] = False,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
@@ -910,7 +957,12 @@ class NodeNMAPPortScanAction(AbstractAction):
"application",
"NMAP",
"port_scan",
{"target_ip_address": target_ip_address, "target_port": target_port, "target_protocol": target_protocol},
{
"target_ip_address": target_ip_address,
"target_port": target_port,
"target_protocol": target_protocol,
"show": show,
},
]
@@ -926,6 +978,7 @@ class NodeNetworkServiceReconAction(AbstractAction):
target_ip_address: Union[str, List[str]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[str, List[str]]] = None,
show: Optional[bool] = False,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
@@ -935,7 +988,12 @@ class NodeNetworkServiceReconAction(AbstractAction):
"application",
"NMAP",
"network_service_recon",
{"target_ip_address": target_ip_address, "target_port": target_port, "target_protocol": target_protocol},
{
"target_ip_address": target_ip_address,
"target_port": target_port,
"target_protocol": target_protocol,
"show": show,
},
]

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

@@ -107,6 +107,8 @@ class AbstractAgent(ABC):
:type observation_space: Optional[ObservationSpace]
:param reward_function: Reward function for the agent.
:type reward_function: Optional[RewardFunction]
:param agent_settings: Configurable Options for Abstracted Agents
:type agent_settings: Optional[AgentSettings]
"""
self.agent_name: str = agent_name or "unnamed_agent"
self.action_manager: Optional[ActionManager] = action_space

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

@@ -5,7 +5,7 @@ from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict
from primaite import getLogger
from primaite import DEFAULT_BANDWIDTH, getLogger
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent
from primaite.game.agent.observations.observation_manager import ObservationManager
@@ -406,7 +406,7 @@ class PrimaiteGame:
for link_cfg in links_cfg:
node_a = net.get_node_by_hostname(link_cfg["endpoint_a_hostname"])
node_b = net.get_node_by_hostname(link_cfg["endpoint_b_hostname"])
bandwidth = link_cfg.get("bandwidth", 100) # default value if not configured
bandwidth = link_cfg.get("bandwidth", DEFAULT_BANDWIDTH) # default value if not configured
if isinstance(node_a, Switch):
endpoint_a = node_a.network_interface[link_cfg["endpoint_a_port"]]

4
src/primaite/interface/request.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Dict, ForwardRef, List, Literal, Union
from pydantic import BaseModel, ConfigDict, StrictBool, validate_call
from pydantic import BaseModel, ConfigDict, StrictBool # , validate_call
RequestFormat = List[Union[str, int, float, Dict]]
@@ -31,7 +31,7 @@ class RequestResponse(BaseModel):
# much. However, in the future we might consider making them mandatory.
@classmethod
@validate_call
# @validate_call # this slows down execution quite a bit.
def from_bool(cls, status_bool: StrictBool) -> RequestResponse:
"""
Construct a basic request response from a boolean.

2
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb
View File

@@ -458,7 +458,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

4
src/primaite/notebooks/Requests-and-Responses.ipynb
View File

@@ -6,9 +6,9 @@
"source": [
"# Requests and Responses\n",
"\n",
"Agents interact with the PrimAITE simulation via the Request system.\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
"Agents interact with the PrimAITE simulation via the Request system.\n"
]
},
{

2
src/primaite/notebooks/Using-Episode-Schedules.ipynb
View File

@@ -423,7 +423,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

2
src/primaite/session/environment.py
View File

@@ -60,7 +60,7 @@ class PrimaiteGymEnv(gymnasium.Env):
next_obs = self._get_obs() # this doesn't update observation, just gets the current observation
reward = self.agent.reward_function.current_reward
_LOGGER.info(f"step: {self.game.step_counter}, Blue reward: {reward}")
_LOGGER.debug(f"step: {self.game.step_counter}, Blue reward: {reward}")
terminated = False
truncated = self.game.calculate_truncated()
info = {

12
src/primaite/setup/reset_demo_notebooks.py
View File

@@ -45,11 +45,7 @@ def run(overwrite_existing: bool = True) -> None:
shutil.copy2(src_fp, dst_fp)
_LOGGER.info(f"Reset example notebook: {dst_fp}")
for src_fp in primaite_root.glob("notebooks/_package_data/*"):
dst_fp = example_notebooks_user_dir / "_package_data" / src_fp.name
if should_copy_file(src_fp, dst_fp, overwrite_existing):
if not Path.exists(example_notebooks_user_dir / "_package_data/"):
Path.mkdir(example_notebooks_user_dir / "_package_data/")
print(dst_fp)
shutil.copy2(src_fp, dst_fp)
_LOGGER.info(f"Copied notebook resource to: {dst_fp}")
src = primaite_root / "notebooks/_package_data/"
dst = example_notebooks_user_dir / "_package_data/"
shutil.copytree(src, dst, dirs_exist_ok=True)

2
src/primaite/simulator/core.py
View File

@@ -221,7 +221,7 @@ class SimComponent(BaseModel):
}
return state
@validate_call
# @validate_call # this slows down execution quite a bit.
def apply_request(self, request: RequestFormat, context: Optional[Dict] = None) -> RequestResponse:
"""
Apply a request to a simulation component. Request data is passed in as a 'namespaced' list of strings.

17
src/primaite/simulator/file_system/file.py
View File

@@ -6,6 +6,8 @@ import json
import warnings
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC, FileSystemItemHealthStatus
from primaite.simulator.file_system.file_type import FileType, get_file_type_from_extension
@@ -203,3 +205,18 @@ class File(FileSystemItemABC):
self.deleted = True
self.sys_log.info(f"File deleted {self.folder_name}/{self.name}")
return True
def show(self, markdown: bool = False):
"""
Prints a table of the file, displaying the file type, name and the file's access count.
:param markdown: Flag indicating if output should be in markdown format.
"""
headers = ["File Name", "File Type", "Number of Accesses"]
table = PrettyTable(headers)
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.name}"
table.add_row([self.name, self.file_type.name, self.num_access])
print(table)

101
src/primaite/simulator/file_system/file_system.py
View File

@@ -67,7 +67,7 @@ class FileSystem(SimComponent):
self._create_manager = RequestManager()
def _create_file_action(request: List[Any], context: Any) -> RequestResponse:
file = self.create_file(folder_name=request[0], file_name=request[1])
file = self.create_file(folder_name=request[0], file_name=request[1], force=request[2])
if not file:
return RequestResponse.from_bool(False)
return RequestResponse(
@@ -75,7 +75,7 @@ class FileSystem(SimComponent):
data={
"file_name": file.name,
"folder_name": file.folder_name,
"file_type": file.file_type,
"file_type": file.file_type.name,
"file_size": file.size,
},
)
@@ -111,9 +111,9 @@ class FileSystem(SimComponent):
data={
"file_name": file.name,
"folder_name": file.folder_name,
"file_type": file.file_type,
"file_type": file.file_type.name,
"file_size": file.size,
"file_status": file.health_status,
"file_status": file.health_status.name,
},
)
return RequestResponse.from_bool(False)
@@ -160,6 +160,21 @@ class FileSystem(SimComponent):
"""
return sum(folder.size for folder in self.folders.values())
def show_num_files(self, markdown: bool = False):
"""
Prints a table showing a host's number of file creations & deletions.
:param markdown: Flag indicating if output should be in markdown format.
"""
headers = ["File creations", "File deletions"]
table = PrettyTable(headers)
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Number of Creations & Deletions"
table.add_row([self.num_file_creations, self.num_file_deletions])
print(table)
def show(self, markdown: bool = False, full: bool = False):
"""
Prints a table of the FileSystem, displaying either just folders or full files.
@@ -167,7 +182,7 @@ class FileSystem(SimComponent):
:param markdown: Flag indicating if output should be in markdown format.
:param full: Flag indicating if to show full files.
"""
headers = ["Folder", "Size", "Deleted"]
headers = ["Folder", "Size", "Health status", "Visible health status", "Deleted"]
if full:
headers[0] = "File Path"
table = PrettyTable(headers)
@@ -178,14 +193,38 @@ class FileSystem(SimComponent):
folders = {**self.folders, **self.deleted_folders}
for folder in folders.values():
if not full:
table.add_row([folder.name, folder.size_str, folder.deleted])
table.add_row(
[
folder.name,
folder.size_str,
folder.health_status.name,
folder.visible_health_status.name,
folder.deleted,
]
)
else:
files = {**folder.files, **folder.deleted_files}
if not files:
table.add_row([folder.name, folder.size_str, folder.deleted])
table.add_row(
[
folder.name,
folder.size_str,
folder.health_status.name,
folder.visible_health_status.name,
folder.deleted,
]
)
else:
for file in files.values():
table.add_row([file.path, file.size_str, file.deleted])
table.add_row(
[
file.path,
file.size_str,
file.health_status.name,
file.visible_health_status.name,
file.deleted,
]
)
if full:
print(table.get_string(sortby="File Path"))
else:
@@ -201,15 +240,15 @@ class FileSystem(SimComponent):
:param folder_name: The name of the folder.
"""
# check if folder with name already exists
if self.get_folder(folder_name):
raise Exception(f"Cannot create folder as it already exists: {folder_name}")
folder = Folder(name=folder_name, sys_log=self.sys_log)
folder = self.get_folder(folder_name)
if folder:
self.sys_log.info(f"Cannot create folder as it already exists: {folder_name}")
else:
folder = Folder(name=folder_name, sys_log=self.sys_log)
self._folder_request_manager.add_request(
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
self.folders[folder.uuid] = folder
self._folder_request_manager.add_request(
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
return folder
def delete_folder(self, folder_name: str) -> bool:
@@ -289,6 +328,7 @@ class FileSystem(SimComponent):
size: Optional[int] = None,
file_type: Optional[FileType] = None,
folder_name: Optional[str] = None,
force: Optional[bool] = False,
) -> File:
"""
Creates a File and adds it to the list of files.
@@ -297,6 +337,7 @@ class FileSystem(SimComponent):
:param size: The size the file takes on disk in bytes.
:param file_type: The type of the file.
:param folder_name: The folder to add the file to.
:param force: Replaces the file if it already exists.
"""
if folder_name:
# check if file with name already exists
@@ -308,17 +349,23 @@ class FileSystem(SimComponent):
# Use root folder if folder_name not supplied
folder = self.get_folder("root")
# Create the file and add it to the folder
file = File(
name=file_name,
sim_size=size,
file_type=file_type,
folder_id=folder.uuid,
folder_name=folder.name,
sim_root=self.sim_root,
sys_log=self.sys_log,
)
folder.add_file(file)
file = self.get_file(folder, file_name)
if file:
self.sys_log.info(f"Cannot create file {file_name} as it already exists.")
if force:
self.sys_log.info(f"Replacing {file_name}")
else:
# Create the file and add it to the folder
file = File(
name=file_name,
sim_size=size,
file_type=file_type,
folder_id=folder.uuid,
folder_name=folder.name,
sim_root=self.sim_root,
sys_log=self.sys_log,
)
folder.add_file(file, force=force)
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
# increment file creation
self.num_file_creations += 1

21
src/primaite/simulator/file_system/file_type.py
View File

@@ -96,6 +96,20 @@ class FileType(Enum):
DB = 32
"Generic DB file. Used by sqlite3."
# Script file types
PS1 = 33
"A Powershell script. Used in Windows operating systems."
BAT = 34
"Windows batch file. Used in Windows operating systems."
SH = 35
"Linux shell script. Used in Unix based operating systems."
# Executable file types
PE = 36
"Portable Executable. Used in Windows operating systems."
ELF = 37
"Executable and Linkable Format. Used in Unix based operating systems."
@classmethod
def _missing_(cls, value: Any) -> FileType:
return cls.UNKNOWN
@@ -116,8 +130,11 @@ class FileType(Enum):
Returns 0 if a default size does not exist.
"""
size = file_type_sizes_bytes[self]
return size if size else 0
try:
size = file_type_sizes_bytes[self]
return size
except KeyError:
return 0
def get_file_type_from_extension(file_type_extension: str) -> FileType:

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

@@ -11,6 +11,7 @@ from typing import Any, Dict, Optional, Type, TypeVar, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, Field
import primaite.simulator.network.nmne
from primaite import getLogger
from primaite.exceptions import NetworkError
from primaite.interface.request import RequestResponse
@@ -1049,7 +1050,7 @@ class Node(SimComponent):
def show_nic(self, markdown: bool = False):
"""Prints a table of the NICs on the Node."""
table = PrettyTable(["Port", "Type", "MAC Address", "Address", "Speed", "Status"])
table = PrettyTable(["Port", "Type", "MAC Address", "Address", "Speed", "Status", "NMNE"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -1066,6 +1067,7 @@ class Node(SimComponent):
ip_address,
network_interface.speed,
"Enabled" if network_interface.enabled else "Disabled",
network_interface.nmne if primaite.simulator.network.nmne.CAPTURE_NMNE else "Disabled",
]
)
print(table)

5
src/primaite/simulator/network/networks.py
View File

@@ -185,7 +185,10 @@ def arcd_uc2_network() -> Network:
db_client_2.run()
web_browser_2 = client_2.software_manager.software.get("WebBrowser")
web_browser_2.target_url = "http://arcd.com/users/"
network.connect(endpoint_b=client_2.network_interface[1], endpoint_a=switch_2.network_interface[2])
network.connect(
endpoint_b=client_2.network_interface[1],
endpoint_a=switch_2.network_interface[2],
)
# Domain Controller
domain_controller = Server(

4
src/primaite/simulator/network/transmission/network_layer.py
View File

@@ -25,6 +25,10 @@ class IPProtocol(Enum):
ICMP = "icmp"
"""Internet Control Message Protocol."""
def model_dump(self) -> str:
"""Return as JSON-serialisable string."""
return self.name
class Precedence(Enum):
"""

4
src/primaite/simulator/network/transmission/transport_layer.py
View File

@@ -72,6 +72,10 @@ class Port(Enum):
POSTGRES_SERVER = 5432
"Postgres SQL Server."
def model_dump(self) -> str:
"""Return a json-serialisable string."""
return self.name
class UDPHeader(BaseModel):
"""

8
src/primaite/simulator/system/applications/nmap.py
View File

@@ -87,7 +87,9 @@ class NMAP(Application):
def _init_request_manager(self) -> RequestManager:
def _ping_scan_action(request: List[Any], context: Any) -> RequestResponse:
results = self.ping_scan(target_ip_address=request[0]["target_ip_address"], json_serializable=True)
results = self.ping_scan(
target_ip_address=request[0]["target_ip_address"], show=request[0]["show"], json_serializable=True
)
if not self._can_perform_network_action():
return RequestResponse.from_bool(False)
return RequestResponse(
@@ -311,7 +313,6 @@ class NMAP(Application):
self.sys_log.info(
f"{self.name}: Responding to port scan request for port {payload.port.value} "
f"({payload.port.name}) over {payload.protocol.name}",
True,
)
self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
@@ -370,7 +371,8 @@ class NMAP(Application):
port_open = self._check_port_open_on_ip_address(ip_address=ip_address, port=port, protocol=protocol)
if port_open:
table.add_row([ip_address, port.value, port.name, protocol.name])
if show:
table.add_row([ip_address, port.value, port.name, protocol.name])
_ip_address = ip_address if not json_serializable else str(ip_address)
_protocol = protocol if not json_serializable else protocol.value
_port = port if not json_serializable else port.value

2
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -143,7 +143,7 @@ class ARP(Service):
self.sys_log.info(f"Cannot send ARP request to a network address {str(target_ip_address)}")
return
if target_ip_address == outbound_network_interface.ip_network.broadcast_address:
self.sys_log.info(f"Cannot send ARP request to a broadcast address {str(target_ip_address)}")
self.sys_log.info(f"Cannot send ARP request to a broadcast addresss {str(target_ip_address)}")
return
self.sys_log.info(f"Sending ARP request from NIC {outbound_network_interface} for ip {target_ip_address}")

6
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -61,7 +61,7 @@ class ICMP(Service):
if target_ip_address.is_loopback:
self.sys_log.info("Pinging loopback address")
return any(network_interface.enabled for network_interface in self.network_interfaces.values())
self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True)
self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=False)
sequence, identifier = 0, None
while sequence < pings:
sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings)
@@ -77,7 +77,7 @@ class ICMP(Service):
f"Received = {request_replies}, "
f"Lost = {pings - request_replies} ({(pings - request_replies) / pings * 100}% loss)"
)
self.sys_log.info(output, to_terminal=True)
self.sys_log.info(output, to_terminal=False)
return passed
@@ -167,7 +167,7 @@ class ICMP(Service):
f"bytes={len(frame.payload)}, "
f"time={time_str}, "
f"TTL={frame.ip.ttl}",
to_terminal=True,
to_terminal=False,
)
if not self.request_replies.get(frame.icmp.identifier):
self.request_replies[frame.icmp.identifier] = 0

1
src/primaite/utils/session_metadata_parser.py
View File

@@ -3,7 +3,6 @@
raise DeprecationWarning(
"Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
)
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import json
from pathlib import Path
from typing import Any, Dict, Union