diff --git a/example_config.yaml b/example_config.yaml index dd5971a1..e3871b4a 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -22,12 +22,12 @@ game_config: actions: - type: DONOTHING nodes: - - ref: client_2 + - node_ref: client_2 actions: - type: LOGON - type: LOGOFF applications: - - ref: client_2_web_browser + - application_ref: client_2_web_browser actions: - type: EXECUTE execution_definition: @@ -44,12 +44,12 @@ game_config: observation_space: network: nodes: - - ref: client_1 + - node_ref: client_1 observations: - logon_status - operating_status services: - - ref: data_manipulation_bot + - service_ref: data_manipulation_bot observations: - operating_status - health_status @@ -59,13 +59,13 @@ game_config: - type: DO_NOTHING network: nodes: - - ref: client_1 + - node_ref: client_1 actions: - type: SCAN - type: LOGON - type: LOGOFF services: - - ref: data_manipulation_bot + - service_ref: data_manipulation_bot actions: - type: COMPROMISE execution_definition: @@ -89,18 +89,18 @@ game_config: observation_space: network: nodes: - - ref: router_1 #TODO: more sub-options here - - ref: switch_1 - - ref: switch_2 - - ref: domain_controller - - ref: web_server - - ref: database_server - - ref: backup_server - - ref: security_suite - - ref: client_1 - - ref: client_2 + - node_ref: router_1 #TODO: more sub-options here + - node_ref: switch_1 + - node_ref: switch_2 + - node_ref: domain_controller + - node_ref: web_server + - node_ref: database_server + - node_ref: backup_server + - node_ref: security_suite + - node_ref: client_1 + - node_ref: client_2 links: - - ref: ... # + - link_ref: ... # acl: ... # ics: ... # @@ -110,7 +110,7 @@ game_config: - type: DO_NOTHING network: nodes: - - ref: router_1 + - node_ref: router_1 reward_function: # ... agent_settings: @@ -153,7 +153,7 @@ simulation: protocol: ICMP - ref: switch_1 - type: swtich + type: switch hostname: switch_1 num_ports: 8 @@ -173,7 +173,7 @@ simulation: type: DNSServer options: domain_mapping: - - arcd.com: 192.168.1.12 # web server + arcd.com: 192.168.1.12 # web server - ref: web_server @@ -254,52 +254,52 @@ simulation: links: - ref: router_1___switch_1 - endpoint_a: router_1 + endpoint_a_ref: router_1 endpoint_a_port: 1 - endpoint_b: switch_1 + endpoint_b_ref: switch_1 endpoint_b_port: 8 - ref: router_1___switch_2 - endpoint_a: router_1 + endpoint_a_ref: router_1 endpoint_a_port: 2 - endpoint_b: switch_2 + endpoint_b_ref: switch_2 endpoint_b_port: 8 - ref: switch_1___domain_controller - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 1 - endpoint_b: domain_controller + endpoint_b_ref: domain_controller endpoint_b_port: 1 - ref: switch_1___web_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 2 - endpoint_b: web_server + endpoint_b_ref: web_server endpoint_b_port: 1 - ref: switch_1___database_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 3 - endpoint_b: database_server + endpoint_b_ref: database_server endpoint_b_port: 1 - ref: switch_1___backup_server - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 4 - endpoint_b: backup_server + endpoint_b_ref: backup_server endpoint_b_port: 1 - ref: switch_1___security_suite - endpoint_a: switch_1 + endpoint_a_ref: switch_1 endpoint_a_port: 7 - endpoint_b: security_suite + endpoint_b_ref: security_suite endpoint_b_port: 1 - ref: switch_2___client_1 - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 1 - endpoint_b: client_1 + endpoint_b_ref: client_1 endpoint_b_port: 1 - ref: switch_2___client_2 - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 2 - endpoint_b: client_2 + endpoint_b_ref: client_2 endpoint_b_port: 1 - ref: switch_2___security_suite - endpoint_a: switch_2 + endpoint_a_ref: switch_2 endpoint_a_port: 7 - endpoint_b: security_suite + endpoint_b_ref: security_suite endpoint_b_port: 2 diff --git a/sandbox.ipynb b/sandbox.ipynb index 5d611ada..aa39c3e9 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -33,28 +33,29 @@ "from primaite.simulator.network.hardware.nodes.router import ACLAction\n", "from primaite.simulator.network.transmission.network_layer import IPProtocol\n", "from primaite.simulator.network.transmission.transport_layer import Port\n", - "\n" + "\n", + "from ipaddress import IPv4Address\n" ] }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-26 10:51:10,388: Added node 48e6cb0b-f351-47f6-b837-df9443f9db26 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,390::WARNING::primaite.simulator.network.container::181::Can't add node 48e6cb0b-f351-47f6-b837-df9443f9db26. It is already in the network.\n", - "2023-09-26 10:51:10,394: Added node 6a969d4d-e0af-402e-b576-2a787505f7c7 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,397: Added node c58e6f17-dbf1-4c6a-9dbf-d60883c6d948 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,401: Added node 7f2a418d-2d0b-4f02-beb3-5703fc5035c8 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,408: Added node 967417fa-2300-4ee1-8ba0-7a4d055d5d30 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,413: Added node 80c1c99b-4c7a-41fb-86f0-b93c35c3b497 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,418: Added node 9a11dd40-9243-4510-9b43-9f247f669ad2 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,424: Added node 81fff4a6-35c8-4933-bb6c-fd8fd49315fe to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n", - "2023-09-26 10:51:10,429: Added node 7cc11532-3f65-4c65-a4df-af2c6318a976 to Network 7250d818-ec1b-4940-bb87-8e05fea87fe9\n" + "2023-09-26 11:47:11,032: Added node bc149bf5-ccc4-4dcd-b419-629ec44b2c9a to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,035: Added node 9cacbaee-33cc-4423-a6c8-fe3dd75b1f87 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,042: Added node d4444d66-7cc3-4cd4-acbd-202cb9fe37ff to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,045: Added node af170371-e99b-42b7-9525-65ca64522539 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,049: Added node d6218f34-a104-469d-a08b-97329ad84c19 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,052: Added node 831a3803-ae65-4cee-a17e-9c1220035bc9 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,055: Added node 1b935654-065d-4cb9-82d9-d67fe3d3304e to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,059: Added node dd181916-076b-4d8a-ab97-a32052624b09 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,064: Added node 3137ab20-1a3c-49f2-8ee5-c862216b2435 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", + "2023-09-26 11:47:11,067: Added node 6ff8b634-7750-4c6d-8109-abf52514dae5 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n" ] }, { @@ -70,6 +71,11 @@ "# import yaml\n", "\n", "\n", + "from typing import Dict\n", + "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", + "from primaite.simulator.system.services.service import Service\n", + "\n", + "\n", "class PrimaiteSession:\n", "\n", " def __init__(self):\n", @@ -78,7 +84,11 @@ "\n", " @classmethod\n", " def from_config(cls, cfg_path):\n", - " ref_to_uuid = {}\n", + " ref_map_nodes: Dict[str,Node] = {}\n", + " ref_map_services: Dict[str, Service] = {}\n", + " ref_map_links: Dict[str, Link] = {}\n", + " # ref_map_agents: Dict[str, AgentInterface] = {}\n", + "\n", "\n", " game = cls()\n", " with open(cfg_path, 'r') as file:\n", @@ -129,6 +139,8 @@ " dst_ip_address = r_cfg.get('ip_address'),\n", " position = r_num\n", " )\n", + " else:\n", + " print('invalid node type')\n", " if 'services' in node_cfg:\n", " for service_cfg in node_cfg['services']:\n", " service_ref = service_cfg['ref']\n", @@ -144,15 +156,43 @@ " }\n", " if service_type in service_types_mapping:\n", " new_node.software_manager.install(service_types_mapping[service_type])\n", - " service_obj = new_node.software_manager.software[service_type]\n", - " ref_to_uuid[service_ref] = service_obj.uuid\n", + " new_service = new_node.software_manager.software[service_type]\n", + " ref_map_services[service_ref] = new_service\n", " else:\n", " print(f\"service type not found {service_type}\")\n", - "\n", + " # service-dependent options\n", + " if service_type == 'DatabaseClient':\n", + " if 'options' in service_cfg:\n", + " opt = service_cfg['options']\n", + " if 'db_server_ip' in opt:\n", + " new_service.configure(server_ip_address=IPv4Address(opt['db_server_ip']))\n", + " if service_type == 'DNSServer':\n", + " if 'options' in service_cfg:\n", + " opt = service_cfg['options']\n", + " if 'domain_mapping' in opt:\n", + " for domain, ip in opt['domain_mapping'].items():\n", + " new_service.dns_register(domain, ip)\n", + " if 'nics' in node_cfg:\n", + " for nic_num, nic_cfg in node_cfg['nics'].items():\n", + " new_node.connect_nic(NIC(ip_address=nic_cfg['ip_address'], subnet_mask=nic_cfg['subnet_mask']))\n", "\n", " net.add_node(new_node)\n", - " ref_to_uuid[node_ref] = new_node.uuid\n", + " ref_map_nodes[node_ref] = new_node.uuid\n", "\n", + " #2. create links between nodes\n", + " for link_cfg in links_cfg:\n", + " node_a = net.nodes[ref_map_nodes[link_cfg['endpoint_a_ref']]]\n", + " node_b = net.nodes[ref_map_nodes[link_cfg['endpoint_b_ref']]]\n", + " if isinstance(node_a, Switch):\n", + " endpoint_a = node_a.switch_ports[link_cfg['endpoint_a_port']]\n", + " else:\n", + " endpoint_a = node_a.ethernet_port[link_cfg['endpoint_a_port']]\n", + " if isinstance(node_b, Switch):\n", + " endpoint_b = node_b.switch_ports[link_cfg['endpoint_b_port']]\n", + " else:\n", + " endpoint_b = node_b.ethernet_port[link_cfg['endpoint_b_port']]\n", + " new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)\n", + " ref_map_links[link_cfg['ref']] = new_link.uuid\n", "\n", " #2. start/setup simulation objects\n", " #3. create agents\n", diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 7ab9b093..66686797 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -221,7 +221,7 @@ class Network(SimComponent): _LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}") self._node_action_manager.remove_action(name=node.uuid) - def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None: + def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> Optional[Link]: """ Connect two endpoints on the network by creating a link between their NICs/SwitchPorts. @@ -248,6 +248,7 @@ class Network(SimComponent): self._nx_graph.add_edge(endpoint_a.parent.hostname, endpoint_b.parent.hostname) link.parent = self _LOGGER.debug(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}") + return link def remove_link(self, link: Link) -> None: """Disconnect a link from the network.