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

2755 Add ability to extend HostNode, NetworkNode, Service and Application outside PrimAITE.

Browse Source
This commit is contained in:
=
2024-09-04 15:49:37 +01:00
parent 08f742b3ec
commit 310876cd3b
17 changed files with 1926 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/primaite/config/load.py
View File

@@ -59,3 +59,17 @@ def data_manipulation_marl_config_path() -> Path:
_LOGGER.error(msg)
raise FileNotFoundError(msg)
return path
def get_extended_config_path() -> Path:
"""
Get the path to an 'extended' example config that contains nodes using the extension framework
:return: Path to the extended example config
:rtype: Path
"""
path = _EXAMPLE_CFG / "extended_config.yaml"
if not path.exists():
msg = f"Example config does not exist: {path}. Have you run `primaite setup`?"
_LOGGER.error(msg)
raise FileNotFoundError(msg)
return path

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

@@ -20,9 +20,10 @@ from primaite.simulator import SIM_OUTPUT
from primaite.simulator.network.airspace import AirSpaceFrequency
from primaite.simulator.network.hardware.base import NetworkInterface, NodeOperatingState, UserManager
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.host_node import NIC, HostNode
from primaite.simulator.network.hardware.nodes.host.server import Printer, Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
@@ -278,8 +279,25 @@ class PrimaiteGame:
for node_cfg in nodes_cfg:
n_type = node_cfg["type"]
new_node = None
if n_type == "computer":
# Handle extended nodes
if n_type.lower() in HostNode._registry:
new_node = HostNode._registry[n_type](
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=IPv4Address(node_cfg.get("subnet_mask", "255.255.255.0")),
default_gateway=node_cfg.get("default_gateway"),
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 in NetworkNode._registry:
new_node = NetworkNode._registry[n_type](
**node_cfg
)
# Default PrimAITE nodes
elif n_type == "computer":
new_node = Computer(
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
@@ -351,10 +369,18 @@ class PrimaiteGame:
for service_cfg in node_cfg["services"]:
new_service = None
service_type = service_cfg["type"]
if service_type in SERVICE_TYPES_MAPPING:
service_class = None
# Handle extended services
if service_type.lower() in Service._registry:
service_class = Service._registry[service_type.lower()]
elif service_type in SERVICE_TYPES_MAPPING:
service_class = SERVICE_TYPES_MAPPING[service_type]
if service_class is not None:
_LOGGER.debug(f"installing {service_type} on node {new_node.hostname}")
new_node.software_manager.install(SERVICE_TYPES_MAPPING[service_type])
new_service = new_node.software_manager.software[service_type]
new_node.software_manager.install(service_class)
new_service = new_node.software_manager.software[service_class.__name__]
# fixing duration for the service
if "fix_duration" in service_cfg.get("options", {}):
@@ -398,8 +424,8 @@ class PrimaiteGame:
new_application = None
application_type = application_cfg["type"]
if application_type in Application._application_registry:
new_node.software_manager.install(Application._application_registry[application_type])
if application_type in Application._registry:
new_node.software_manager.install(Application._registry[application_type])
new_application = new_node.software_manager.software[application_type] # grab the instance
# fixing duration for the application

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

@@ -12,7 +12,9 @@ from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.network.airspace import AirSpace
from primaite.simulator.network.hardware.base import Link, Node, WiredNetworkInterface
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.network.hardware.nodes.host.server import Printer
from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service
@@ -128,6 +130,16 @@ class Network(SimComponent):
def firewall_nodes(self) -> List[Node]:
"""The Firewalls in the Network."""
return [node for node in self.nodes.values() if node.__class__.__name__ == "Firewall"]
@property
def extended_hostnodes(self) -> List[Node]:
"""Extended nodes that inherited HostNode in the network"""
return [node for node in self.nodes.values() if node.__class__.__name__.lower() in HostNode._registry]
@property
def extended_networknodes(self) -> List[Node]:
"""Extended nodes that inherited NetworkNode in the network"""
return [node for node in self.nodes.values() if node.__class__.__name__.lower() in NetworkNode._registry]
@property
def printer_nodes(self) -> List[Node]:
@@ -160,6 +172,7 @@ class Network(SimComponent):
"Printer": self.printer_nodes,
"Wireless Router": self.wireless_router_nodes,
}
if nodes:
table = PrettyTable(["Node", "Type", "Operating State"])
if markdown:

2
src/primaite/simulator/network/hardware/base.py
View File

@@ -1699,7 +1699,7 @@ class Node(SimComponent):
if self.software_manager.software.get(application_name):
self.sys_log.warning(f"Can't install {application_name}. It's already installed.")
return RequestResponse(status="success", data={"reason": "already installed"})
application_class = Application._application_registry[application_name]
application_class = Application._registry[application_name]
self.software_manager.install(application_class)
application_instance = self.software_manager.software.get(application_name)
self.applications[application_instance.uuid] = application_instance

22
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -2,7 +2,7 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Optional
from typing import Any, ClassVar, Dict, Optional, Type
from primaite import getLogger
from primaite.simulator.network.hardware.base import (
@@ -325,10 +325,30 @@ class HostNode(Node):
network_interface: Dict[int, NIC] = {}
"The NICs on the node by port id."
_registry: ClassVar[Dict[str, Type["HostNode"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init__(self, ip_address: IPV4Address, subnet_mask: IPV4Address, **kwargs):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
def __init_subclass__(cls, identifier: str = 'default', **kwargs: Any) -> None:
"""
Register a hostnode type.
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type identifier: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
"""
if identifier == 'default':
return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower()
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
@property
def nmap(self) -> Optional[NMAP]:
"""

21
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import Optional
from typing import Any, ClassVar, Dict, Optional, Type
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
@@ -16,6 +16,25 @@ class NetworkNode(Node):
provide functionality for receiving and processing frames received on their network interfaces.
"""
_registry: ClassVar[Dict[str, Type["NetworkNode"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: str = 'default', **kwargs: Any) -> None:
"""
Register a networknode type.
:param identifier: Uniquely specifies an networknode class by name. Used for finding items by config.
:type identifier: str
:raises ValueError: When attempting to register an networknode with a name that is already allocated.
"""
if identifier == 'default':
return
identifier = identifier.lower()
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Tried to define new networknode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
@abstractmethod
def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface):
"""

10
src/primaite/simulator/system/applications/application.py
View File

@@ -41,10 +41,10 @@ class Application(IOSoftware):
install_countdown: Optional[int] = None
"The countdown to the end of the installation process. None if not currently installing"
_application_registry: ClassVar[Dict[str, Type["Application"]]] = {}
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
def __init_subclass__(cls, identifier: str = 'default', **kwargs: Any) -> None:
"""
Register an application type.
@@ -52,10 +52,12 @@ class Application(IOSoftware):
:type identifier: str
:raises ValueError: When attempting to register an application with a name that is already allocated.
"""
if identifier == 'default':
return
super().__init_subclass__(**kwargs)
if identifier in cls._application_registry:
if identifier in cls._registry:
raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.")
cls._application_registry[identifier] = cls
cls._registry[identifier] = cls
def __init__(self, **kwargs):
super().__init__(**kwargs)

22
src/primaite/simulator/system/services/service.py
View File

@@ -3,7 +3,7 @@ from __future__ import annotations
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, Optional
from typing import Any, ClassVar, Dict, Optional, Type
from primaite import getLogger
from primaite.interface.request import RequestFormat, RequestResponse
@@ -46,9 +46,29 @@ class Service(IOSoftware):
restart_countdown: Optional[int] = None
"If currently restarting, how many timesteps remain until the restart is finished."
_registry: ClassVar[Dict[str, Type["Service"]]] = {}
"""Registry of service types. Automatically populated when subclasses are defined."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
def __init_subclass__(cls, identifier: str = 'default', **kwargs: Any) -> None:
"""
Register a hostnode type.
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type identifier: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
"""
if identifier == 'default':
return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower()
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
def _can_perform_action(self) -> bool:
"""
Checks if the service can perform actions.