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

yaml parse and connect links

Browse Source
This commit is contained in:
Marek Wolan
2023-09-26 11:48:22 +01:00
parent fdf66ba3de
commit 92e0110e73
3 changed files with 99 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
example_config.yaml
View File

@@ -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

74
sandbox.ipynb
View File

@@ -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",

3
src/primaite/simulator/network/container.py
View File

@@ -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.