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

bugfix - Make node schemas stricter

Browse Source
This commit is contained in:
Marek Wolan
2025-02-05 15:04:41 +00:00
parent 4a472c5c75
commit c1abbfe58c
17 changed files with 73 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1529,9 +1529,11 @@ class Node(SimComponent, ABC):
class ConfigSchema(BaseModel, ABC):
"""Configuration Schema for Node based classes."""
model_config = ConfigDict(arbitrary_types_allowed=True)
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
"""Configure pydantic to allow arbitrary types, let the instance have attributes not present in the model."""
type: str
hostname: str
"The node hostname on the network."
@@ -1570,6 +1572,8 @@ class Node(SimComponent, ABC):
operating_state: Any = None
users: Any = None # Temporary to appease "extra=forbid"
config: ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema())
"""Configuration items within Node"""

3
src/primaite/simulator/network/hardware/nodes/host/computer.py
View File

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import ClassVar, Dict
from typing import ClassVar, Dict, Literal
from pydantic import Field
@@ -40,6 +40,7 @@ class Computer(HostNode, discriminator="computer"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Computer class."""
type: Literal["computer"] = "computer"
hostname: str = "Computer"
config: ConfigSchema = Field(default_factory=lambda: Computer.ConfigSchema())

7
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, Literal, Optional
from pydantic import Field
@@ -333,9 +333,14 @@ class HostNode(Node, discriminator="host-node"):
class ConfigSchema(Node.ConfigSchema):
"""Configuration Schema for HostNode class."""
type: Literal["host-node"]
hostname: str = "HostNode"
subnet_mask: IPV4Address = "255.255.255.0"
ip_address: IPV4Address
services: Any = None # temporarily unset to appease extra="forbid"
applications: Any = None # temporarily unset to appease extra="forbid"
folders: Any = None # temporarily unset to appease extra="forbid"
network_interfaces: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema())

4
src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -1,5 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Literal
from pydantic import Field
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
@@ -36,6 +38,7 @@ class Server(HostNode, discriminator="server"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Server class."""
type: Literal["server"] = "server"
hostname: str = "server"
config: ConfigSchema = Field(default_factory=lambda: Server.ConfigSchema())
@@ -49,6 +52,7 @@ class Printer(HostNode, discriminator="printer"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Printer class."""
type: Literal["printer"] = "printer"
hostname: str = "printer"
config: ConfigSchema = Field(default_factory=lambda: Printer.ConfigSchema())

3
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Dict, Final, Union
from typing import Dict, Final, Literal, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -103,6 +103,7 @@ class Firewall(Router, discriminator="firewall"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for Firewall 'Nodes' within PrimAITE."""
type: Literal["firewall"] = "firewall"
hostname: str = "firewall"
num_ports: int = 0

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

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import Optional
from typing import Any, Optional
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
@@ -16,6 +16,11 @@ class NetworkNode(Node, discriminator="network-node"):
provide functionality for receiving and processing frames received on their network interfaces.
"""
class ConfigSchema(Node.ConfigSchema):
"""Config schema for Node baseclass."""
num_ports: Any = None # temporarily unset to appease extra="forbid"
@abstractmethod
def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface):
"""

27
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -4,7 +4,7 @@ from __future__ import annotations
import secrets
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -1204,8 +1204,13 @@ class Router(NetworkNode, discriminator="router"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Routers."""
type: Literal["router"] = "router"
hostname: str = "router"
num_ports: int = 5
acl: Any = None # temporarily unset to appease extra="forbid"
routes: Any = None # temporarily unset to appease extra="forbid"
ports: Any = None # temporarily unset to appease extra="forbid"
default_route: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: Router.ConfigSchema())
@@ -1625,16 +1630,20 @@ class Router(NetworkNode, discriminator="router"):
:return: Configured router.
:rtype: Router
"""
ports = config.pop("ports", None)
acl = config.pop("acl", None)
routes = config.pop("routes", None)
default_route = config.pop("default_route", None)
router = Router(config=Router.ConfigSchema(**config))
if "ports" in config:
for port_num, port_cfg in config["ports"].items():
if ports:
for port_num, port_cfg in ports.items():
router.configure_port(
port=port_num,
ip_address=port_cfg["ip_address"],
subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")),
)
if "acl" in config:
for r_num, r_cfg in config["acl"].items():
if acl:
for r_num, r_cfg in acl.items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
@@ -1646,16 +1655,16 @@ class Router(NetworkNode, discriminator="router"):
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
if "routes" in config:
for route in config.get("routes"):
if routes:
for route in routes:
router.route_table.add_route(
address=IPv4Address(route.get("address")),
subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")),
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
metric=float(route.get("metric", 0)),
)
if "default_route" in config:
next_hop_ip_address = config["default_route"].get("next_hop_ip_address", None)
if default_route:
next_hop_ip_address = default_route.get("next_hop_ip_address", None)
if next_hop_ip_address:
router.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address)
router.operating_state = (

3
src/primaite/simulator/network/hardware/nodes/network/switch.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Optional
from typing import Dict, Literal, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
@@ -101,6 +101,7 @@ class Switch(NetworkNode, discriminator="switch"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Switch nodes within PrimAITE."""
type: Literal["switch"] = "switch"
hostname: str = "Switch"
num_ports: int = 24
"The number of ports on the switch. Default is 24."

13
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Literal, Optional, Union
from pydantic import Field, validate_call
@@ -126,10 +126,13 @@ class WirelessRouter(Router, discriminator="wireless-router"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for WirelessRouter nodes within PrimAITE."""
type: Literal["wireless-router"] = "wireless-router"
hostname: str = "WirelessRouter"
airspace: AirSpace
num_ports: int = 0
router_interface: Any = None # temporarily unset to appease extra="forbid"
wireless_access_point: Any = None # temporarily unset to appease extra="forbid"
airspace: AirSpace
config: ConfigSchema = Field(default_factory=lambda: WirelessRouter.ConfigSchema())
def __init__(self, **kwargs):
@@ -137,7 +140,7 @@ class WirelessRouter(Router, discriminator="wireless-router"):
self.connect_nic(
WirelessAccessPoint(
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=kwargs["config"].airspace
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=self.airspace
)
)
@@ -236,7 +239,7 @@ class WirelessRouter(Router, discriminator="wireless-router"):
)
@classmethod
def from_config(cls, config: Dict, **kwargs) -> "WirelessRouter":
def from_config(cls, config: Dict, airspace: AirSpace) -> "WirelessRouter":
"""Generate the wireless router from config.
Schema:
@@ -263,7 +266,7 @@ class WirelessRouter(Router, discriminator="wireless-router"):
:return: WirelessRouter instance.
:rtype: WirelessRouter
"""
router = cls(config=cls.ConfigSchema(**config))
router = cls(config=cls.ConfigSchema(**config), airspace=airspace)
router.operating_state = (
NodeOperatingState.ON if not (p := config.get("operating_state")) else NodeOperatingState[p.upper()]
)