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

Merge remote-tracking branch 'origin/dev' into feature/2319-database-admin

Browse Source
This commit is contained in:
Marek Wolan
2024-03-01 09:23:00 +00:00
parent bd0b2e0033 b849bfd6e2
commit 9ff8adab1a
88 changed files with 3595 additions and 668 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -572,7 +572,7 @@ class NetworkNICDisableAction(NetworkNICAbstractAction):
class ActionManager:
"""Class which manages the action space for an agent."""
_act_class_identifiers: Dict[str, type] = {
act_class_identifiers: Dict[str, type] = {
"DONOTHING": DoNothingAction,
"NODE_SERVICE_SCAN": NodeServiceScanAction,
"NODE_SERVICE_STOP": NodeServiceStopAction,
@@ -733,7 +733,7 @@ class ActionManager:
# and `options` is an optional dict of options to pass to the init method of the action class
act_type = act_spec.get("type")
act_options = act_spec.get("options", {})
self.actions[act_type] = self._act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.actions[act_type] = self.act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.action_map: Dict[int, Tuple[str, Dict]] = {}
"""

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

@@ -13,7 +13,7 @@ the structure:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
node_name: database_server
folder_name: database
file_name: database.db
@@ -21,7 +21,7 @@ the structure:
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
node_name: web_server
service_ref: web_server_database_client
```
"""
@@ -184,7 +184,7 @@ class WebServer404Penalty(AbstractReward):
service_name = config.get("service_name")
if not (node_hostname and service_name):
msg = (
f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not "
f"{cls.__name__} could not be initialised from config because node_name and service_ref were not "
"found in reward config."
)
_LOGGER.warning(msg)
@@ -285,13 +285,14 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
class RewardFunction:
"""Manages the reward function for the agent."""
__rew_class_identifiers: Dict[str, Type[AbstractReward]] = {
rew_class_identifiers: Dict[str, Type[AbstractReward]] = {
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty,
"WEBPAGE_UNAVAILABLE_PENALTY": WebpageUnavailablePenalty,
"GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY": GreenAdminDatabaseUnreachablePenalty,
}
"""List of reward class identifiers."""
def __init__(self):
"""Initialise the reward function object."""
@@ -338,7 +339,7 @@ class RewardFunction:
for rew_component_cfg in config["reward_components"]:
rew_type = rew_component_cfg["type"]
weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type]
rew_class = cls.rew_class_identifiers[rew_type]
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}))
new.register_component(component=rew_instance, weight=weight)
return new

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

@@ -16,6 +16,7 @@ from primaite.simulator.network.hardware.base import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.nmne import set_nmne_config
@@ -42,6 +43,7 @@ APPLICATION_TYPES_MAPPING = {
"DataManipulationBot": DataManipulationBot,
"DoSBot": DoSBot,
}
"""List of available applications that can be installed on nodes in the PrimAITE Simulation."""
SERVICE_TYPES_MAPPING = {
"DNSClient": DNSClient,
@@ -53,6 +55,7 @@ SERVICE_TYPES_MAPPING = {
"NTPClient": NTPClient,
"NTPServer": NTPServer,
}
"""List of available services that can be installed on nodes in the PrimAITE Simulation."""
class PrimaiteGameOptions(BaseModel):
@@ -229,28 +232,36 @@ class PrimaiteGame:
new_node = Computer(
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=node_cfg["subnet_mask"],
subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")),
default_gateway=node_cfg["default_gateway"],
dns_server=node_cfg["dns_server"],
operating_state=NodeOperatingState.ON,
dns_server=node_cfg.get("dns_server", None),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "server":
new_node = Server(
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=node_cfg["subnet_mask"],
subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")),
default_gateway=node_cfg["default_gateway"],
dns_server=node_cfg.get("dns_server"),
operating_state=NodeOperatingState.ON,
dns_server=node_cfg.get("dns_server", None),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "switch":
new_node = Switch(
hostname=node_cfg["hostname"],
num_ports=node_cfg.get("num_ports"),
operating_state=NodeOperatingState.ON,
num_ports=int(node_cfg.get("num_ports", "8")),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
elif n_type == "router":
new_node = Router.from_config(node_cfg)
elif n_type == "firewall":
new_node = Firewall.from_config(node_cfg)
else:
_LOGGER.warning(f"invalid node type {n_type} in config")
if "services" in node_cfg:
@@ -263,6 +274,9 @@ class PrimaiteGame:
new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service.uuid
# start the service
new_service.start()
else:
msg = f"Configuration contains an invalid service type: {service_type}"
_LOGGER.error(msg)
@@ -282,18 +296,16 @@ class PrimaiteGame:
if service_type == "DatabaseService":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.password = opt.get("backup_server_ip", None)
new_service.configure_backup(backup_server=IPv4Address(opt.get("backup_server_ip")))
new_service.start()
if service_type == "FTPServer":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.server_password = opt.get("server_password")
new_service.start()
if service_type == "NTPClient":
if "options" in service_cfg:
opt = service_cfg["options"]
new_service.ntp_server = IPv4Address(opt.get("ntp_server_ip"))
new_service.start()
if "applications" in node_cfg:
for application_cfg in node_cfg["applications"]:
new_application = None
@@ -309,13 +321,16 @@ class PrimaiteGame:
_LOGGER.error(msg)
raise ValueError(msg)
# run the application
new_application.run()
if application_type == "DataManipulationBot":
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.configure(
server_ip_address=IPv4Address(opt.get("server_ip")),
server_password=opt.get("server_password"),
payload=opt.get("payload"),
payload=opt.get("payload", "DELETE"),
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
data_manipulation_p_of_success=float(opt.get("data_manipulation_p_of_success", "0.1")),
)
@@ -330,7 +345,6 @@ class PrimaiteGame:
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.target_url = opt.get("target_url")
elif application_type == "DoSBot":
if "options" in application_cfg:
opt = application_cfg["options"]
@@ -347,10 +361,20 @@ class PrimaiteGame:
for nic_num, nic_cfg in node_cfg["network_interfaces"].items():
new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"]))
# temporarily set to 0 so all nodes are initially on
new_node.start_up_duration = 0
new_node.shut_down_duration = 0
net.add_node(new_node)
new_node.power_on()
# run through the power on step if the node is to be turned on at the start
if new_node.operating_state == NodeOperatingState.ON:
new_node.power_on()
game.ref_map_nodes[node_ref] = new_node.uuid
# set start up and shut down duration
new_node.start_up_duration = int(node_cfg.get("start_up_duration", 3))
new_node.shut_down_duration = int(node_cfg.get("shut_down_duration", 3))
# 2. create links between nodes
for link_cfg in links_cfg:
node_a = net.nodes[game.ref_map_nodes[link_cfg["endpoint_a_ref"]]]