diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index b79fc985..e22c882c 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -219,6 +219,45 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction): self.verb: str = "fix" +class NodeApplicationInstallAction(AbstractAction): + """Action which installs an application.""" + + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"node_id": num_nodes} + + def form_request(self, node_id: int, application_name: str, ip_address: str) -> 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: + return ["do_nothing"] + return [ + "network", + "node", + node_name, + "software_manager", + "application", + "install", + application_name, + ip_address, + ] + + +class NodeApplicationRemoveAction(AbstractAction): + """Action which removes/uninstalls an application.""" + + def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: + super().__init__(manager=manager) + self.shape: Dict[str, int] = {"node_id": num_nodes} + + def form_request(self, node_id: int, application_name: str) -> 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: + return ["do_nothing"] + return ["network", "node", node_name, "software_manager", "application", "uninstall", application_name] + + class NodeFolderAbstractAction(AbstractAction): """ Base class for folder actions. @@ -658,6 +697,8 @@ class ActionManager: "NODE_APPLICATION_SCAN": NodeApplicationScanAction, "NODE_APPLICATION_CLOSE": NodeApplicationCloseAction, "NODE_APPLICATION_FIX": NodeApplicationFixAction, + "NODE_APPLICATION_INSTALL": NodeApplicationInstallAction, + "NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction, "NODE_FILE_SCAN": NodeFileScanAction, "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, "NODE_FILE_DELETE": NodeFileDeleteAction, diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 38d20e1f..bfa547d2 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -5,7 +5,7 @@ import secrets from abc import ABC, abstractmethod from ipaddress import IPv4Address, IPv4Network from pathlib import Path -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional, Type, TypeVar, Union from prettytable import MARKDOWN, PrettyTable from pydantic import BaseModel, Field @@ -35,8 +35,11 @@ from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.processes.process import Process from primaite.simulator.system.services.service import Service +from primaite.simulator.system.software import IOSoftware from primaite.utils.validators import IPV4Address +IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware) + _LOGGER = getLogger(__name__) @@ -843,12 +846,58 @@ class Node(SimComponent): ) rm.add_request("os", RequestType(func=self._os_request_manager, validator=_node_is_on)) + self._software_request_manager = RequestManager() + rm.add_request("software_manager", RequestType(func=self._software_request_manager, validator=_node_is_on)) + self._application_manager = RequestManager() + self._software_request_manager.add_request( + name="application", request_type=RequestType(func=self._application_manager) + ) + + self._application_manager.add_request( + name="install", + request_type=RequestType( + func=lambda request, context: RequestResponse.from_bool( + self.application_install_action( + application=self._read_application_type(request[0]), ip_address=request[1] + ) + ) + ), + ) + + self._application_manager.add_request( + name="uninstall", + request_type=RequestType( + func=lambda request, context: RequestResponse.from_bool( + self.application_uninstall_action(application=self._read_application_type(request[0])) + ) + ), + ) + return rm def _install_system_software(self): """Install System Software - software that is usually provided with the OS.""" pass + def _read_application_type(self, application_class_str: str) -> Type[IOSoftwareClass]: + """Wrapper that converts the string from the request manager into the appropriate class for the application.""" + if application_class_str == "DoSBot": + from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot + + return DoSBot + elif application_class_str == "DataManipulationBot": + from primaite.simulator.system.applications.red_applications.data_manipulation_bot import ( + DataManipulationBot, + ) + + return DataManipulationBot + elif application_class_str == "WebBrowser": + from primaite.simulator.system.applications.web_browser import WebBrowser + + return WebBrowser + else: + return 0 + def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -1257,6 +1306,75 @@ class Node(SimComponent): _LOGGER.info(f"Removed application {application.name} from node {self.hostname}") self._application_request_manager.remove_request(application.name) + def application_install_action(self, application: Application, ip_address: Optional[str] = None) -> bool: + """ + Install an application on this node and configure it. + + This method is useful for allowing agents to take this action. + + :param application: Application object that has not been installed on any node yet. + :type application: Application + :param ip_address: IP address used to configure the application + (target IP for the DoSBot or server IP for the DataManipulationBot) + :type ip_address: str + :return: True if the application is installed successfully, otherwise False. + """ + if application in self: + _LOGGER.warning( + f"Can't add application {application.__name__}" + f"to node {self.hostname}. It's already installed." + ) + return True + + self.software_manager.install(application) + application_instance = self.software_manager.software.get(str(application.__name__)) + self.applications[application_instance.uuid] = application_instance + self.sys_log.info(f"Installed application {application_instance.name}") + _LOGGER.debug(f"Added application {application_instance.name} to node {self.hostname}") + self._application_request_manager.add_request( + application_instance.name, RequestType(func=application_instance._request_manager) + ) + + # Configure application if additional parameters are given + if ip_address: + if application_instance.name == "DoSBot": + application_instance.configure(target_ip_address=IPv4Address(ip_address)) + elif application_instance.name == "DataManipulationBot": + application_instance.configure(server_ip_address=IPv4Address(ip_address)) + else: + pass + + if application_instance.name in self.software_manager.software: + return True + else: + return False + + def application_uninstall_action(self, application: Application) -> bool: + """ + Uninstall and completely remove application from this node. + + This method is useful for allowing agents to take this action. + + :param application: Application object that is currently associated with this node. + :type application: Application + :return: True if the application is uninstalled successfully, otherwise False. + """ + if application.__name__ not in self.software_manager.software: + _LOGGER.warning( + f"Can't remove application {application.__name__}" + f"from node {self.hostname}. It's not installed." + ) + return True + + application_instance = self.software_manager.software.get( + str(application.__name__) + ) # This works because we can't have two applications with the same name on the same node + # self.uninstall_application(application_instance) + self.software_manager.uninstall(application_instance.name) + + if application_instance.name not in self.software_manager.software: + return True + else: + return False + def _shut_down_actions(self): """Actions to perform when the node is shut down.""" # Turn off all the services in the node @@ -1288,4 +1406,6 @@ class Node(SimComponent): def __contains__(self, item: Any) -> bool: if isinstance(item, Service): return item.uuid in self.services + elif isinstance(item, Application): + return item.uuid in self.applications return None diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index ab60adde..3ab32bc6 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -88,7 +88,7 @@ class Software(SimComponent): "The count of times the software has been scanned, defaults to 0." revealed_to_red: bool = False "Indicates if the software has been revealed to red agent, defaults is False." - software_manager: "SoftwareManager" = None + software_manager: Optional["SoftwareManager"] = None "An instance of Software Manager that is used by the parent node." sys_log: SysLog = None "An instance of SysLog that is used by the parent node." diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml new file mode 100644 index 00000000..c1908fc4 --- /dev/null +++ b/tests/assets/configs/test_application_install.yaml @@ -0,0 +1,986 @@ +training_config: + rl_framework: SB3 + rl_algorithm: PPO + seed: 333 + n_learn_episodes: 1 + n_eval_episodes: 5 + max_steps_per_episode: 128 + deterministic_eval: false + n_agents: 1 + agent_references: + - defender + +io_settings: + save_agent_actions: true + save_step_metadata: false + save_pcap_logs: false + save_sys_logs: false + + +game: + max_episode_length: 128 + ports: + - HTTP + - POSTGRES_SERVER + protocols: + - ICMP + - TCP + - UDP + thresholds: + nmne: + high: 10 + medium: 5 + low: 0 + +agents: + - ref: client_2_green_user + team: GREEN + type: ProbabilisticAgent + agent_settings: + action_probabilities: + 0: 0.3 + 1: 0.6 + 2: 0.1 + observation_space: + type: UC2GreenObservation + action_space: + action_list: + - type: DONOTHING + - type: NODE_APPLICATION_EXECUTE + options: + nodes: + - node_name: client_2 + applications: + - application_name: WebBrowser + - application_name: DatabaseClient + max_folders_per_node: 1 + max_files_per_folder: 1 + max_services_per_node: 1 + max_applications_per_node: 2 + action_map: + 0: + action: DONOTHING + options: {} + 1: + action: NODE_APPLICATION_EXECUTE + options: + node_id: 0 + application_id: 0 + 2: + action: NODE_APPLICATION_EXECUTE + options: + node_id: 0 + application_id: 1 + + reward_function: + reward_components: + - type: WEBPAGE_UNAVAILABLE_PENALTY + weight: 0.25 + options: + node_hostname: client_2 + - type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY + weight: 0.05 + options: + node_hostname: client_2 + + - ref: client_1_green_user + team: GREEN + type: ProbabilisticAgent + agent_settings: + action_probabilities: + 0: 0.3 + 1: 0.6 + 2: 0.1 + observation_space: + type: UC2GreenObservation + action_space: + action_list: + - type: DONOTHING + - type: NODE_APPLICATION_EXECUTE + options: + nodes: + - node_name: client_1 + applications: + - application_name: WebBrowser + - application_name: DatabaseClient + max_folders_per_node: 1 + max_files_per_folder: 1 + max_services_per_node: 1 + max_applications_per_node: 2 + action_map: + 0: + action: DONOTHING + options: {} + 1: + action: NODE_APPLICATION_EXECUTE + options: + node_id: 0 + application_id: 0 + 2: + action: NODE_APPLICATION_EXECUTE + options: + node_id: 0 + application_id: 1 + + reward_function: + reward_components: + - type: WEBPAGE_UNAVAILABLE_PENALTY + weight: 0.25 + options: + node_hostname: client_1 + - type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY + weight: 0.05 + options: + node_hostname: client_1 + + + + + + - ref: data_manipulation_attacker + team: RED + type: RedDatabaseCorruptingAgent + + observation_space: + type: UC2RedObservation + options: + nodes: {} + + action_space: + action_list: + - type: DONOTHING + - type: NODE_APPLICATION_EXECUTE + options: + nodes: + - node_name: client_1 + applications: + - application_name: DataManipulationBot + - node_name: client_2 + applications: + - application_name: DataManipulationBot + max_folders_per_node: 1 + max_files_per_folder: 1 + max_services_per_node: 1 + + reward_function: + reward_components: + - type: DUMMY + + agent_settings: # options specific to this particular agent type, basically args of __init__(self) + start_settings: + start_step: 25 + frequency: 20 + variance: 5 + + - ref: defender + team: BLUE + type: ProxyAgent + + observation_space: + type: UC2BlueObservation + options: + num_services_per_node: 1 + num_folders_per_node: 1 + num_files_per_folder: 1 + num_nics_per_node: 2 + nodes: + - node_hostname: domain_controller + services: + - service_name: DNSServer + - node_hostname: web_server + services: + - service_name: WebServer + - node_hostname: database_server + folders: + - folder_name: database + files: + - file_name: database.db + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 + links: + - link_ref: router_1___switch_1 + - link_ref: router_1___switch_2 + - link_ref: switch_1___domain_controller + - link_ref: switch_1___web_server + - link_ref: switch_1___database_server + - link_ref: switch_1___backup_server + - link_ref: switch_1___security_suite + - link_ref: switch_2___client_1 + - link_ref: switch_2___client_2 + - link_ref: switch_2___security_suite + acl: + options: + max_acl_rules: 10 + router_hostname: router_1 + ip_address_order: + - node_hostname: domain_controller + nic_num: 1 + - node_hostname: web_server + nic_num: 1 + - node_hostname: database_server + nic_num: 1 + - node_hostname: backup_server + nic_num: 1 + - node_hostname: security_suite + nic_num: 1 + - node_hostname: client_1 + nic_num: 1 + - node_hostname: client_2 + nic_num: 1 + - node_hostname: security_suite + nic_num: 2 + ics: null + + action_space: + action_list: + - type: DONOTHING + - type: NODE_SERVICE_SCAN + - type: NODE_SERVICE_STOP + - type: NODE_SERVICE_START + - type: NODE_SERVICE_PAUSE + - type: NODE_SERVICE_RESUME + - type: NODE_SERVICE_RESTART + - type: NODE_SERVICE_DISABLE + - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_FIX + - type: NODE_FILE_SCAN + - type: NODE_FILE_CHECKHASH + - type: NODE_FILE_DELETE + - type: NODE_FILE_REPAIR + - type: NODE_FILE_RESTORE + - type: NODE_FOLDER_SCAN + - type: NODE_FOLDER_CHECKHASH + - type: NODE_FOLDER_REPAIR + - type: NODE_FOLDER_RESTORE + - type: NODE_OS_SCAN + - type: NODE_SHUTDOWN + - type: NODE_STARTUP + - type: NODE_RESET + - type: NETWORK_ACL_ADDRULE + options: + target_router_hostname: router_1 + - type: NETWORK_ACL_REMOVERULE + options: + target_router_hostname: router_1 + - type: NETWORK_NIC_ENABLE + - type: NETWORK_NIC_DISABLE + - type: NODE_APPLICATION_INSTALL + - type: NODE_APPLICATION_REMOVE + - type: NODE_APPLICATION_EXECUTE + + action_map: + 0: + action: DONOTHING + options: {} + # scan webapp service + 1: + action: NODE_SERVICE_SCAN + options: + node_id: 1 + service_id: 0 + # stop webapp service + 2: + action: NODE_SERVICE_STOP + options: + node_id: 1 + service_id: 0 + # start webapp service + 3: + action: "NODE_SERVICE_START" + options: + node_id: 1 + service_id: 0 + 4: + action: "NODE_SERVICE_PAUSE" + options: + node_id: 1 + service_id: 0 + 5: + action: "NODE_SERVICE_RESUME" + options: + node_id: 1 + service_id: 0 + 6: + action: "NODE_SERVICE_RESTART" + options: + node_id: 1 + service_id: 0 + 7: + action: "NODE_SERVICE_DISABLE" + options: + node_id: 1 + service_id: 0 + 8: + action: "NODE_SERVICE_ENABLE" + options: + node_id: 1 + service_id: 0 + 9: # check database.db file + action: "NODE_FILE_SCAN" + options: + node_id: 2 + folder_id: 0 + file_id: 0 + 10: + action: "NODE_FILE_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + options: + node_id: 2 + folder_id: 0 + file_id: 0 + 11: + action: "NODE_FILE_DELETE" + options: + node_id: 2 + folder_id: 0 + file_id: 0 + 12: + action: "NODE_FILE_REPAIR" + options: + node_id: 2 + folder_id: 0 + file_id: 0 + 13: + action: "NODE_SERVICE_FIX" + options: + node_id: 2 + service_id: 0 + 14: + action: "NODE_FOLDER_SCAN" + options: + node_id: 2 + folder_id: 0 + 15: + action: "NODE_FOLDER_SCAN" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context. + options: + node_id: 2 + folder_id: 0 + 16: + action: "NODE_FOLDER_REPAIR" + options: + node_id: 2 + folder_id: 0 + 17: + action: "NODE_FOLDER_RESTORE" + options: + node_id: 2 + folder_id: 0 + 18: + action: "NODE_OS_SCAN" + options: + node_id: 0 + 19: + action: "NODE_SHUTDOWN" + options: + node_id: 0 + 20: + action: NODE_STARTUP + options: + node_id: 0 + 21: + action: NODE_RESET + options: + node_id: 0 + 22: + action: "NODE_OS_SCAN" + options: + node_id: 1 + 23: + action: "NODE_SHUTDOWN" + options: + node_id: 1 + 24: + action: NODE_STARTUP + options: + node_id: 1 + 25: + action: NODE_RESET + options: + node_id: 1 + 26: # old action num: 18 + action: "NODE_OS_SCAN" + options: + node_id: 2 + 27: + action: "NODE_SHUTDOWN" + options: + node_id: 2 + 28: + action: NODE_STARTUP + options: + node_id: 2 + 29: + action: NODE_RESET + options: + node_id: 2 + 30: + action: "NODE_OS_SCAN" + options: + node_id: 3 + 31: + action: "NODE_SHUTDOWN" + options: + node_id: 3 + 32: + action: NODE_STARTUP + options: + node_id: 3 + 33: + action: NODE_RESET + options: + node_id: 3 + 34: + action: "NODE_OS_SCAN" + options: + node_id: 4 + 35: + action: "NODE_SHUTDOWN" + options: + node_id: 4 + 36: + action: NODE_STARTUP + options: + node_id: 4 + 37: + action: NODE_RESET + options: + node_id: 4 + 38: + action: "NODE_OS_SCAN" + options: + node_id: 5 + 39: # old action num: 19 # shutdown client 1 + action: "NODE_SHUTDOWN" + options: + node_id: 5 + 40: # old action num: 20 + action: NODE_STARTUP + options: + node_id: 5 + 41: # old action num: 21 + action: NODE_RESET + options: + node_id: 5 + 42: + action: "NODE_OS_SCAN" + options: + node_id: 6 + 43: + action: "NODE_SHUTDOWN" + options: + node_id: 6 + 44: + action: NODE_STARTUP + options: + node_id: 6 + 45: + action: NODE_RESET + options: + node_id: 6 + + 46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1" + action: "NETWORK_ACL_ADDRULE" + options: + position: 1 + permission: 2 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL + source_port_id: 1 + dest_port_id: 1 + protocol_id: 1 + 47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2" + action: "NETWORK_ACL_ADDRULE" + options: + position: 2 + permission: 2 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL + source_port_id: 1 + dest_port_id: 1 + protocol_id: 1 + 48: # old action num: 24 # block tcp traffic from client 1 to web app + action: "NETWORK_ACL_ADDRULE" + options: + position: 3 + permission: 2 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 + 49: # old action num: 25 # block tcp traffic from client 2 to web app + action: "NETWORK_ACL_ADDRULE" + options: + position: 4 + permission: 2 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 + 50: # old action num: 26 + action: "NETWORK_ACL_ADDRULE" + options: + position: 5 + permission: 2 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 + 51: # old action num: 27 + action: "NETWORK_ACL_ADDRULE" + options: + position: 6 + permission: 2 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database + source_port_id: 1 + dest_port_id: 1 + protocol_id: 3 + 52: # old action num: 28 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 0 + 53: # old action num: 29 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 1 + 54: # old action num: 30 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 2 + 55: # old action num: 31 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 3 + 56: # old action num: 32 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 4 + 57: # old action num: 33 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 5 + 58: # old action num: 34 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 6 + 59: # old action num: 35 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 7 + 60: # old action num: 36 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 8 + 61: # old action num: 37 + action: "NETWORK_ACL_REMOVERULE" + options: + position: 9 + 62: # old action num: 38 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 0 + nic_id: 0 + 63: # old action num: 39 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 0 + nic_id: 0 + 64: # old action num: 40 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 1 + nic_id: 0 + 65: # old action num: 41 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 1 + nic_id: 0 + 66: # old action num: 42 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 2 + nic_id: 0 + 67: # old action num: 43 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 2 + nic_id: 0 + 68: # old action num: 44 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 3 + nic_id: 0 + 69: # old action num: 45 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 3 + nic_id: 0 + 70: # old action num: 46 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 4 + nic_id: 0 + 71: # old action num: 47 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 4 + nic_id: 0 + 72: # old action num: 48 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 4 + nic_id: 1 + 73: # old action num: 49 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 4 + nic_id: 1 + 74: # old action num: 50 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 5 + nic_id: 0 + 75: # old action num: 51 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 5 + nic_id: 0 + 76: # old action num: 52 + action: "NETWORK_NIC_DISABLE" + options: + node_id: 6 + nic_id: 0 + 77: # old action num: 53 + action: "NETWORK_NIC_ENABLE" + options: + node_id: 6 + nic_id: 0 + 78: + action: NODE_APPLICATION_INSTALL + options: + node_id: 0 + application_name: DoSBot + ip_address: 192.168.1.14 + 79: + action: NODE_APPLICATION_REMOVE + options: + node_id: 0 + application_name: DoSBot + 80: + action: NODE_APPLICATION_REMOVE + options: + node_id: 0 + application_name: WebBrowser + 81: + action: NODE_APPLICATION_EXECUTE + options: + node_id: 0 + application_id: 0 + + + + options: + nodes: + - node_name: domain_controller + applications: + - application_name: DoSBot + - node_name: web_server + applications: + - application_name: DatabaseClient + services: + - service_name: WebServer + - node_name: database_server + folders: + - folder_name: database + files: + - file_name: database.db + services: + - service_name: DatabaseService + - node_name: backup_server + - node_name: security_suite + - node_name: client_1 + - node_name: client_2 + + max_folders_per_node: 2 + max_files_per_folder: 2 + max_services_per_node: 2 + max_nics_per_node: 8 + max_acl_rules: 10 + ip_address_order: + - node_name: domain_controller + nic_num: 1 + - node_name: web_server + nic_num: 1 + - node_name: database_server + nic_num: 1 + - node_name: backup_server + nic_num: 1 + - node_name: security_suite + nic_num: 1 + - node_name: client_1 + nic_num: 1 + - node_name: client_2 + nic_num: 1 + - node_name: security_suite + nic_num: 2 + + + reward_function: + reward_components: + - type: DATABASE_FILE_INTEGRITY + weight: 0.40 + options: + node_hostname: database_server + folder_name: database + file_name: database.db + - type: SHARED_REWARD + weight: 1.0 + options: + agent_name: client_1_green_user + - type: SHARED_REWARD + weight: 1.0 + options: + agent_name: client_2_green_user + + + + agent_settings: + flatten_obs: true + + + + + +simulation: + network: + nmne_config: + capture_nmne: true + nmne_capture_keywords: + - DELETE + nodes: + + - ref: router_1 + hostname: router_1 + type: router + num_ports: 5 + ports: + 1: + ip_address: 192.168.1.1 + subnet_mask: 255.255.255.0 + 2: + ip_address: 192.168.10.1 + subnet_mask: 255.255.255.0 + acl: + 18: + action: PERMIT + src_port: POSTGRES_SERVER + dst_port: POSTGRES_SERVER + 19: + action: PERMIT + src_port: DNS + dst_port: DNS + 20: + action: PERMIT + src_port: FTP + dst_port: FTP + 21: + action: PERMIT + src_port: HTTP + dst_port: HTTP + 22: + action: PERMIT + src_port: ARP + dst_port: ARP + 23: + action: PERMIT + protocol: ICMP + + - ref: switch_1 + hostname: switch_1 + type: switch + num_ports: 8 + + - ref: switch_2 + hostname: switch_2 + type: switch + num_ports: 8 + + - ref: domain_controller + hostname: domain_controller + type: server + ip_address: 192.168.1.10 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + services: + - ref: domain_controller_dns_server + type: DNSServer + options: + domain_mapping: + arcd.com: 192.168.1.12 # web server + + - ref: web_server + hostname: web_server + type: server + ip_address: 192.168.1.12 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + services: + - ref: web_server_web_service + type: WebServer + applications: + - ref: web_server_database_client + type: DatabaseClient + options: + db_server_ip: 192.168.1.14 + + + - ref: database_server + hostname: database_server + type: server + ip_address: 192.168.1.14 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + services: + - ref: database_service + type: DatabaseService + options: + backup_server_ip: 192.168.1.16 + - ref: database_ftp_client + type: FTPClient + + - ref: backup_server + hostname: backup_server + type: server + ip_address: 192.168.1.16 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + services: + - ref: backup_service + type: FTPServer + + - ref: security_suite + hostname: security_suite + type: server + ip_address: 192.168.1.110 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + dns_server: 192.168.1.10 + network_interfaces: + 2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot + ip_address: 192.168.10.110 + subnet_mask: 255.255.255.0 + + - ref: client_1 + hostname: client_1 + type: computer + ip_address: 192.168.10.21 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.10.1 + dns_server: 192.168.1.10 + applications: + - ref: data_manipulation_bot + type: DataManipulationBot + options: + port_scan_p_of_success: 0.8 + data_manipulation_p_of_success: 0.8 + payload: "DELETE" + server_ip: 192.168.1.14 + - ref: client_1_web_browser + type: WebBrowser + options: + target_url: http://arcd.com/users/ + - ref: client_1_database_client + type: DatabaseClient + options: + db_server_ip: 192.168.1.14 + services: + - ref: client_1_dns_client + type: DNSClient + + - ref: client_2 + hostname: client_2 + type: computer + ip_address: 192.168.10.22 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.10.1 + dns_server: 192.168.1.10 + applications: + - ref: client_2_web_browser + type: WebBrowser + options: + target_url: http://arcd.com/users/ + - ref: data_manipulation_bot + type: DataManipulationBot + options: + port_scan_p_of_success: 0.8 + data_manipulation_p_of_success: 0.8 + payload: "DELETE" + server_ip: 192.168.1.14 + - ref: client_2_database_client + type: DatabaseClient + options: + db_server_ip: 192.168.1.14 + services: + - ref: client_2_dns_client + type: DNSClient + + + + links: + - ref: router_1___switch_1 + endpoint_a_ref: router_1 + endpoint_a_port: 1 + endpoint_b_ref: switch_1 + endpoint_b_port: 8 + - ref: router_1___switch_2 + endpoint_a_ref: router_1 + endpoint_a_port: 2 + endpoint_b_ref: switch_2 + endpoint_b_port: 8 + - ref: switch_1___domain_controller + endpoint_a_ref: switch_1 + endpoint_a_port: 1 + endpoint_b_ref: domain_controller + endpoint_b_port: 1 + - ref: switch_1___web_server + endpoint_a_ref: switch_1 + endpoint_a_port: 2 + endpoint_b_ref: web_server + endpoint_b_port: 1 + - ref: switch_1___database_server + endpoint_a_ref: switch_1 + endpoint_a_port: 3 + endpoint_b_ref: database_server + endpoint_b_port: 1 + - ref: switch_1___backup_server + endpoint_a_ref: switch_1 + endpoint_a_port: 4 + endpoint_b_ref: backup_server + endpoint_b_port: 1 + - ref: switch_1___security_suite + endpoint_a_ref: switch_1 + endpoint_a_port: 7 + endpoint_b_ref: security_suite + endpoint_b_port: 1 + - ref: switch_2___client_1 + endpoint_a_ref: switch_2 + endpoint_a_port: 1 + endpoint_b_ref: client_1 + endpoint_b_port: 1 + - ref: switch_2___client_2 + endpoint_a_ref: switch_2 + endpoint_a_port: 2 + endpoint_b_ref: client_2 + endpoint_b_port: 1 + - ref: switch_2___security_suite + endpoint_a_ref: switch_2 + endpoint_a_port: 7 + endpoint_b_ref: security_suite + endpoint_b_port: 2 diff --git a/tests/conftest.py b/tests/conftest.py index 078a78bd..fd8727cb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -480,6 +480,8 @@ def game_and_agent(): {"type": "NODE_APPLICATION_SCAN"}, {"type": "NODE_APPLICATION_CLOSE"}, {"type": "NODE_APPLICATION_FIX"}, + {"type": "NODE_APPLICATION_INSTALL"}, + {"type": "NODE_APPLICATION_REMOVE"}, {"type": "NODE_FILE_SCAN"}, {"type": "NODE_FILE_CHECKHASH"}, {"type": "NODE_FILE_DELETE"}, @@ -507,10 +509,16 @@ def game_and_agent(): nodes=[ { "node_name": "client_1", - "applications": [{"application_name": "WebBrowser"}], + "applications": [ + {"application_name": "WebBrowser"}, + {"application_name": "DoSBot"}, + ], "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}], }, - {"node_name": "server_1", "services": [{"service_name": "DNSServer"}]}, + { + "node_name": "server_1", + "services": [{"service_name": "DNSServer"}], + }, {"node_name": "server_2", "services": [{"service_name": "WebServer"}]}, {"node_name": "router"}, ], diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index b68a887e..0b31a353 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -1,8 +1,13 @@ +import yaml + +from primaite.game.game import PrimaiteGame +from primaite.session.environment import PrimaiteGymEnv from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.network.hardware.nodes.host.server import Server from primaite.simulator.system.applications.database_client import DatabaseClient from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot from primaite.simulator.system.services.database.database_service import DatabaseService +from tests import TEST_ASSETS_ROOT def test_data_manipulation(uc2_network): @@ -32,3 +37,41 @@ def test_data_manipulation(uc2_network): # Now check that the DB client on the web_server can successfully query the users table on the database assert db_client.query("SELECT") + + +def test_application_install_uninstall_on_uc2(): + """Test Application install and uninstall via agent actions mid episode.""" + with open(TEST_ASSETS_ROOT / "configs/test_application_install.yaml", "r") as f: + cfg = yaml.safe_load(f) + + env = PrimaiteGymEnv(game_config=cfg) + env.agent.flatten_obs = False + env.reset() + + _, _, _, _, _ = env.step(0) + domcon = env.game.simulation.network.get_node_by_hostname("domain_controller") + + # Test we cannot execute the DoSBot app as it is not installed yet + _, _, _, _, info = env.step(81) + assert info["agent_actions"]["defender"].response.status == "unreachable" + + # Test we can Install the DoSBot app + _, _, _, _, info = env.step(78) + assert "DoSBot" in domcon.software_manager.software + + # Test we can now execute the DoSBot app + _, _, _, _, info = env.step(81) + assert info["agent_actions"]["defender"].response.status == "success" + + # Test we can Uninstall the DoSBot app + _, _, _, _, info = env.step(79) + assert "DoSBot" not in domcon.software_manager.software + + # Test we cannot execute the DoSBot app as it was uninstalled + _, _, _, _, info = env.step(81) + assert info["agent_actions"]["defender"].response.status == "unreachable" + + # Test we can uninstall one of the default apps (WebBrowser) + assert "WebBrowser" in domcon.software_manager.software + _, _, _, _, info = env.step(80) + assert "WebBrowser" not in domcon.software_manager.software diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index b3a52cd8..96e3f390 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -10,6 +10,7 @@ # 4. Check that the simulation has changed in the way that I expect. # 5. Repeat for all actions. +from ipaddress import IPv4Address from typing import Tuple import pytest @@ -455,3 +456,28 @@ def test_node_application_close_integration(game_and_agent: Tuple[PrimaiteGame, game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED + + +def test_node_application_install_and_uninstall_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the NodeApplicationInstallAction and NodeApplicationRemoveAction can form a request and that + it is accepted by the simulation. + + When you initiate a install action, the Application will be installed and configured on the node. + The remove action will uninstall the application from the node.""" + game, agent = game_and_agent + + client_1 = game.simulation.network.get_node_by_hostname("client_1") + + assert client_1.software_manager.software.get("DoSBot") is None + + action = ("NODE_APPLICATION_INSTALL", {"node_id": 0, "application_name": "DoSBot", "ip_address": "192.168.1.14"}) + agent.store_action(action) + game.step() + + assert client_1.software_manager.software.get("DoSBot") is not None + + action = ("NODE_APPLICATION_REMOVE", {"node_id": 0, "application_name": "DoSBot"}) + agent.store_action(action) + game.step() + + assert client_1.software_manager.software.get("DoSBot") is None