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

Add methods for adding/removing nodes form network

Browse Source
This commit is contained in:
Marek Wolan
2023-08-23 14:41:30 +01:00
parent 7e64acd368
commit 1613bbe27a
3 changed files with 93 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

200
src/primaite/notebooks/create-simulation.ipynb
View File

@@ -18,7 +18,7 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -36,26 +36,12 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'uuid': '5304ed6d-de4c-408c-ae24-ada32852d196',\n",
" 'network': {'uuid': 'fa17dfe8-81a1-4c7f-8c5b-8c2d3b1e8756',\n",
" 'nodes': {},\n",
" 'links': {}},\n",
" 'domain': {'uuid': '320cbb83-eb1b-4911-a4f0-fc46d8038a8a', 'accounts': {}}}"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"my_sim = Simulation()\n",
"net = my_sim.network\n",
"my_sim.describe_state()"
]
},
@@ -68,7 +54,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -77,17 +63,14 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"my_pc = Node(hostname=\"primaite_pc\",)\n",
"net.add_node(my_pc)\n",
"my_server = Node(hostname=\"google_server\")\n",
"\n",
"# TODO: when there is a proper function for adding nodes, use it instead of manually adding.\n",
"\n",
"my_sim.network.nodes[my_pc.uuid] = my_pc\n",
"my_sim.network.nodes[my_server.uuid] = my_server\n"
"net.add_node(my_server)\n"
]
},
{
@@ -99,7 +82,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -108,22 +91,12 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2023-08-20 18:42:51,310: NIC 5c:b6:26:c0:86:61/130.1.1.1 connected to Link 5c:b6:26:c0:86:61/130.1.1.1<-->01:ef:b1:a3:24:72\n",
"2023-08-20 18:42:51,311: SwitchPort 01:ef:b1:a3:24:72 connected to Link 5c:b6:26:c0:86:61/130.1.1.1<-->01:ef:b1:a3:24:72\n",
"2023-08-20 18:42:51,314: NIC f6:de:1e:63:8e:7f/130.1.1.2 connected to Link f6:de:1e:63:8e:7f/130.1.1.2<-->30:9e:c8:d4:5d:f3\n",
"2023-08-20 18:42:51,315: SwitchPort 30:9e:c8:d4:5d:f3 connected to Link f6:de:1e:63:8e:7f/130.1.1.2<-->30:9e:c8:d4:5d:f3\n"
]
}
],
"outputs": [],
"source": [
"my_swtich = Switch(hostname=\"switch1\", num_ports=12)\n",
"net.add_node(my_swtich)\n",
"\n",
"pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
"my_pc.connect_nic(pc_nic)\n",
@@ -149,7 +122,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -159,7 +132,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -169,20 +142,9 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"FileSystemFile(uuid='253e4606-0f6d-4e57-8db0-6fa7e331ecea', name='favicon.ico', size=40.0, file_type=<FileSystemFileType.PNG: '11'>, action_manager=None)"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"my_server_folder = my_server.file_system.create_folder(\"static\")\n",
"my_server.file_system.create_file(\"favicon.ico\", file_type=FileSystemFileType.PNG)"
@@ -197,7 +159,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -213,7 +175,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -222,7 +184,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -238,7 +200,7 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -247,7 +209,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -264,130 +226,18 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'uuid': '5304ed6d-de4c-408c-ae24-ada32852d196',\n",
" 'network': {'uuid': 'fa17dfe8-81a1-4c7f-8c5b-8c2d3b1e8756',\n",
" 'nodes': {'1fa46446-6681-4e25-a3ba-c4c2cc564630': {'uuid': '1fa46446-6681-4e25-a3ba-c4c2cc564630',\n",
" 'hostname': 'primaite_pc',\n",
" 'operating_state': 0,\n",
" 'NICs': {'09ca02eb-7733-492c-9eff-f0d6b6ebeeda': {'uuid': '09ca02eb-7733-492c-9eff-f0d6b6ebeeda',\n",
" 'ip_adress': '130.1.1.1',\n",
" 'subnet_mask': '255.255.255.0',\n",
" 'gateway': '130.1.1.255',\n",
" 'mac_address': '5c:b6:26:c0:86:61',\n",
" 'speed': 100,\n",
" 'mtu': 1500,\n",
" 'wake_on_lan': False,\n",
" 'dns_servers': [],\n",
" 'enabled': False}},\n",
" 'file_system': {'uuid': '8b533e31-04e9-4838-839d-0656ace3e57a',\n",
" 'folders': {'b450c223-872c-4fe0-90cc-9da80973eaad': {'uuid': 'b450c223-872c-4fe0-90cc-9da80973eaad',\n",
" 'name': 'downloads',\n",
" 'size': 1000.0,\n",
" 'files': {'8160e685-a76f-4171-8a12-3d6b32a9ea16': {'uuid': '8160e685-a76f-4171-8a12-3d6b32a9ea16',\n",
" 'name': 'firefox_installer.zip',\n",
" 'size': 1000.0,\n",
" 'file_type': 'ZIP'}},\n",
" 'is_quarantined': False}}},\n",
" 'applications': {'c82f1064-f35e-466b-88ae-3f61ba0e5161': {'uuid': 'c82f1064-f35e-466b-88ae-3f61ba0e5161',\n",
" 'health_state': 'GOOD',\n",
" 'health_state_red_view': 'GOOD',\n",
" 'criticality': 'MEDIUM',\n",
" 'patching_count': 0,\n",
" 'scanning_count': 0,\n",
" 'revealed_to_red': False,\n",
" 'installing_count': 0,\n",
" 'max_sessions': 1,\n",
" 'tcp': True,\n",
" 'udp': True,\n",
" 'ports': ['HTTP'],\n",
" 'opearting_state': 'RUNNING',\n",
" 'execution_control_status': 'manual',\n",
" 'num_executions': 0,\n",
" 'groups': []}},\n",
" 'services': {},\n",
" 'process': {}},\n",
" '7f637689-6f91-4026-a685-48a9067f03e8': {'uuid': '7f637689-6f91-4026-a685-48a9067f03e8',\n",
" 'hostname': 'google_server',\n",
" 'operating_state': 0,\n",
" 'NICs': {'1abc7272-c516-4463-bd07-1a3cefe39313': {'uuid': '1abc7272-c516-4463-bd07-1a3cefe39313',\n",
" 'ip_adress': '130.1.1.2',\n",
" 'subnet_mask': '255.255.255.0',\n",
" 'gateway': '130.1.1.255',\n",
" 'mac_address': 'f6:de:1e:63:8e:7f',\n",
" 'speed': 100,\n",
" 'mtu': 1500,\n",
" 'wake_on_lan': False,\n",
" 'dns_servers': [],\n",
" 'enabled': False}},\n",
" 'file_system': {'uuid': 'ac9a6643-8349-4f7a-98c7-a1a9f97ce123',\n",
" 'folders': {'befa5d92-0878-4da2-9dac-f993c0b4a554': {'uuid': 'befa5d92-0878-4da2-9dac-f993c0b4a554',\n",
" 'name': 'static',\n",
" 'size': 0,\n",
" 'files': {},\n",
" 'is_quarantined': False},\n",
" '27383b5e-8884-4ec0-bb50-a5d43e460dfa': {'uuid': '27383b5e-8884-4ec0-bb50-a5d43e460dfa',\n",
" 'name': 'root',\n",
" 'size': 40.0,\n",
" 'files': {'253e4606-0f6d-4e57-8db0-6fa7e331ecea': {'uuid': '253e4606-0f6d-4e57-8db0-6fa7e331ecea',\n",
" 'name': 'favicon.ico',\n",
" 'size': 40.0,\n",
" 'file_type': 'PNG'}},\n",
" 'is_quarantined': False}}},\n",
" 'applications': {},\n",
" 'services': {},\n",
" 'process': {}}},\n",
" 'links': {'a449b1ff-50d9-4342-861e-44f2d4dfef37': {'uuid': 'a449b1ff-50d9-4342-861e-44f2d4dfef37',\n",
" 'endpoint_a': '09ca02eb-7733-492c-9eff-f0d6b6ebeeda',\n",
" 'endpoint_b': 'ee4557d9-a309-45dd-a6e0-5b572cc70ee5',\n",
" 'bandwidth': 100.0,\n",
" 'current_load': 0.0},\n",
" 'ebd7687b-ec69-4f1b-b2ba-86669aa95723': {'uuid': 'ebd7687b-ec69-4f1b-b2ba-86669aa95723',\n",
" 'endpoint_a': '1abc7272-c516-4463-bd07-1a3cefe39313',\n",
" 'endpoint_b': 'dc26b764-a07e-486a-99a4-798c8e0c187a',\n",
" 'bandwidth': 100.0,\n",
" 'current_load': 0.0}}},\n",
" 'domain': {'uuid': '320cbb83-eb1b-4911-a4f0-fc46d8038a8a',\n",
" 'accounts': {'5fdcfb66-84f3-4f0f-a3a7-d0cb0e1a5d51': {'uuid': '5fdcfb66-84f3-4f0f-a3a7-d0cb0e1a5d51',\n",
" 'num_logons': 0,\n",
" 'num_logoffs': 0,\n",
" 'num_group_changes': 0,\n",
" 'username': 'admin',\n",
" 'password': 'admin12',\n",
" 'account_type': 'USER',\n",
" 'enabled': True}}}}"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"my_sim.describe_state()"
]
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'{\"uuid\": \"5304ed6d-de4c-408c-ae24-ada32852d196\", \"network\": {\"uuid\": \"fa17dfe8-81a1-4c7f-8c5b-8c2d3b1e8756\", \"nodes\": {\"1fa46446-6681-4e25-a3ba-c4c2cc564630\": {\"uuid\": \"1fa46446-6681-4e25-a3ba-c4c2cc564630\", \"hostname\": \"primaite_pc\", \"operating_state\": 0, \"NICs\": {\"09ca02eb-7733-492c-9eff-f0d6b6ebeeda\": {\"uuid\": \"09ca02eb-7733-492c-9eff-f0d6b6ebeeda\", \"ip_adress\": \"130.1.1.1\", \"subnet_mask\": \"255.255.255.0\", \"gateway\": \"130.1.1.255\", \"mac_address\": \"5c:b6:26:c0:86:61\", \"speed\": 100, \"mtu\": 1500, \"wake_on_lan\": false, \"dns_servers\": [], \"enabled\": false}}, \"file_system\": {\"uuid\": \"8b533e31-04e9-4838-839d-0656ace3e57a\", \"folders\": {\"b450c223-872c-4fe0-90cc-9da80973eaad\": {\"uuid\": \"b450c223-872c-4fe0-90cc-9da80973eaad\", \"name\": \"downloads\", \"size\": 1000.0, \"files\": {\"8160e685-a76f-4171-8a12-3d6b32a9ea16\": {\"uuid\": \"8160e685-a76f-4171-8a12-3d6b32a9ea16\", \"name\": \"firefox_installer.zip\", \"size\": 1000.0, \"file_type\": \"ZIP\"}}, \"is_quarantined\": false}}}, \"applications\": {\"c82f1064-f35e-466b-88ae-3f61ba0e5161\": {\"uuid\": \"c82f1064-f35e-466b-88ae-3f61ba0e5161\", \"health_state\": \"GOOD\", \"health_state_red_view\": \"GOOD\", \"criticality\": \"MEDIUM\", \"patching_count\": 0, \"scanning_count\": 0, \"revealed_to_red\": false, \"installing_count\": 0, \"max_sessions\": 1, \"tcp\": true, \"udp\": true, \"ports\": [\"HTTP\"], \"opearting_state\": \"RUNNING\", \"execution_control_status\": \"manual\", \"num_executions\": 0, \"groups\": []}}, \"services\": {}, \"process\": {}}, \"7f637689-6f91-4026-a685-48a9067f03e8\": {\"uuid\": \"7f637689-6f91-4026-a685-48a9067f03e8\", \"hostname\": \"google_server\", \"operating_state\": 0, \"NICs\": {\"1abc7272-c516-4463-bd07-1a3cefe39313\": {\"uuid\": \"1abc7272-c516-4463-bd07-1a3cefe39313\", \"ip_adress\": \"130.1.1.2\", \"subnet_mask\": \"255.255.255.0\", \"gateway\": \"130.1.1.255\", \"mac_address\": \"f6:de:1e:63:8e:7f\", \"speed\": 100, \"mtu\": 1500, \"wake_on_lan\": false, \"dns_servers\": [], \"enabled\": false}}, \"file_system\": {\"uuid\": \"ac9a6643-8349-4f7a-98c7-a1a9f97ce123\", \"folders\": {\"befa5d92-0878-4da2-9dac-f993c0b4a554\": {\"uuid\": \"befa5d92-0878-4da2-9dac-f993c0b4a554\", \"name\": \"static\", \"size\": 0, \"files\": {}, \"is_quarantined\": false}, \"27383b5e-8884-4ec0-bb50-a5d43e460dfa\": {\"uuid\": \"27383b5e-8884-4ec0-bb50-a5d43e460dfa\", \"name\": \"root\", \"size\": 40.0, \"files\": {\"253e4606-0f6d-4e57-8db0-6fa7e331ecea\": {\"uuid\": \"253e4606-0f6d-4e57-8db0-6fa7e331ecea\", \"name\": \"favicon.ico\", \"size\": 40.0, \"file_type\": \"PNG\"}}, \"is_quarantined\": false}}}, \"applications\": {}, \"services\": {}, \"process\": {}}}, \"links\": {\"a449b1ff-50d9-4342-861e-44f2d4dfef37\": {\"uuid\": \"a449b1ff-50d9-4342-861e-44f2d4dfef37\", \"endpoint_a\": \"09ca02eb-7733-492c-9eff-f0d6b6ebeeda\", \"endpoint_b\": \"ee4557d9-a309-45dd-a6e0-5b572cc70ee5\", \"bandwidth\": 100.0, \"current_load\": 0.0}, \"ebd7687b-ec69-4f1b-b2ba-86669aa95723\": {\"uuid\": \"ebd7687b-ec69-4f1b-b2ba-86669aa95723\", \"endpoint_a\": \"1abc7272-c516-4463-bd07-1a3cefe39313\", \"endpoint_b\": \"dc26b764-a07e-486a-99a4-798c8e0c187a\", \"bandwidth\": 100.0, \"current_load\": 0.0}}}, \"domain\": {\"uuid\": \"320cbb83-eb1b-4911-a4f0-fc46d8038a8a\", \"accounts\": {\"5fdcfb66-84f3-4f0f-a3a7-d0cb0e1a5d51\": {\"uuid\": \"5fdcfb66-84f3-4f0f-a3a7-d0cb0e1a5d51\", \"num_logons\": 0, \"num_logoffs\": 0, \"num_group_changes\": 0, \"username\": \"admin\", \"password\": \"admin12\", \"account_type\": \"USER\", \"enabled\": true}}}}'"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"import json\n",
"json.dumps(my_sim.describe_state())"

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

@@ -137,6 +137,7 @@ class SimComponent(BaseModel):
kwargs["uuid"] = str(uuid4())
super().__init__(**kwargs)
self.action_manager: Optional[ActionManager] = None
self._parent: Optional["SimComponent"] = None
@abstractmethod
def describe_state(self) -> Dict:
@@ -187,3 +188,24 @@ class SimComponent(BaseModel):
Override this method with anything that needs to happen within the component for it to be reset.
"""
pass
@property
def parent(self) -> "SimComponent":
"""Reference to the parent object which manages this object.
:return: Parent object.
:rtype: SimComponent
"""
return self._parent
@parent.setter
def parent(self, new_parent: "SimComponent") -> None:
if self._parent:
msg = f"Overwriting parent of {self}, {self._parent} with {new_parent}"
_LOGGER.warn(msg)
raise RuntimeWarning(msg)
self._parent = new_parent
@parent.deleter
def parent(self) -> None:
self._parent = None

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

@@ -1,8 +1,11 @@
from typing import Dict
from typing import Any, Dict
from primaite import getLogger
from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent
from primaite.simulator.network.hardware.base import Link, Node
_LOGGER = getLogger(__name__)
class NetworkContainer(SimComponent):
"""Top level container object representing the physical network."""
@@ -40,3 +43,45 @@ class NetworkContainer(SimComponent):
}
)
return state
def add_node(self, node: Node) -> None:
"""
Add an existing node to the network.
:param node: Node instance that the network should keep track of.
:type node: Node
"""
if node in self:
_LOGGER.warning(f"Can't add node {node}. It is already in the network.")
self.nodes[node.uuid] = node
node.parent = self
def remove_node(self, node: Node) -> None:
"""
Remove a node from the network.
:param node: Node instance that is currently part of the network that should be removed.
:type node: Node
"""
if node not in self:
_LOGGER.warning(f"Can't remove node {node}. It's not in the network.")
del self.nodes[node.uuid]
del node.parent # misleading?
def connect_nodes(self, node1: Node, node2: Node) -> None:
"""TODO."""
# I think we should not be forcing users to add and remove individual links.
# Clearly if a link exists between two nodes in the network, then the link is also part of the network.
# I'm just not sure how we ought to handle link creation as it requires an unoccupied network device on the node
raise NotImplementedError
def disconnect_nodes(self, node1: Node, node2: Node) -> None:
"""TODO."""
raise NotImplementedError
def __contains__(self, item: Any) -> bool:
if isinstance(item, Node):
return item.uuid in self.nodes
elif isinstance(item, Link):
return item.uuid in self.links
raise TypeError("")