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

#2887 - End of day commit. Updates to ConfigSchema inheritance, and some initials changes to Router to remove the custom from_config method

Browse Source
This commit is contained in:
Charlie Crane
2025-01-15 16:33:11 +00:00
parent 582e7cfec7
commit 70d9fe2fd9
9 changed files with 46 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -277,67 +277,6 @@ class PrimaiteGame:
if n_type in Node._registry: if n_type in Node._registry:
# simplify down Node creation: # simplify down Node creation:
new_node = Node._registry["n_type"].from_config(config=node_config) new_node = Node._registry["n_type"].from_config(config=node_config)
# Default PrimAITE nodes
# if n_type == "computer":
# new_node = Computer(
# 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 == "server":
# new_node = Server(
# 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 == "switch":
# new_node = Switch(
# hostname=node_cfg["hostname"],
# num_ports=int(node_cfg.get("num_ports", "8")),
# operating_state=NodeOperatingState.ON
# if not (p := node_cfg.get("operating_state"))
# else NodeOperatingState[p.upper()],
# )
# elif n_type == "router":
# new_node = Router.from_config(node_cfg)
# elif n_type == "firewall":
# new_node = Firewall.from_config(node_cfg)
# elif n_type == "wireless_router":
# new_node = WirelessRouter.from_config(node_cfg, airspace=net.airspace)
# elif n_type == "printer":
# new_node = Printer(
# hostname=node_cfg["hostname"],
# ip_address=node_cfg["ip_address"],
# subnet_mask=node_cfg["subnet_mask"],
# operating_state=NodeOperatingState.ON
# if not (p := node_cfg.get("operating_state"))
# else NodeOperatingState[p.upper()],
# )
# # Handle extended nodes
# elif n_type.lower() in Node._registry:
# new_node = HostNode._registry[n_type](
# hostname=node_cfg["hostname"],
# ip_address=node_cfg.get("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)
else: else:
msg = f"invalid node type {n_type} in config" msg = f"invalid node type {n_type} in config"
_LOGGER.error(msg) _LOGGER.error(msg)

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

@@ -9,7 +9,7 @@ from pathlib import Path
from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union
from prettytable import MARKDOWN, PrettyTable from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, Field, validate_call from pydantic import BaseModel, ConfigDict, Field, validate_call
from primaite import getLogger from primaite import getLogger
from primaite.exceptions import NetworkError from primaite.exceptions import NetworkError
@@ -1480,8 +1480,6 @@ class Node(SimComponent, ABC):
:param operating_state: The node operating state, either ON or OFF. :param operating_state: The node operating state, either ON or OFF.
""" """
hostname: str
"The node hostname on the network."
default_gateway: Optional[IPV4Address] = None default_gateway: Optional[IPV4Address] = None
"The default gateway IP address for forwarding network traffic to other networks." "The default gateway IP address for forwarding network traffic to other networks."
operating_state: NodeOperatingState = NodeOperatingState.OFF operating_state: NodeOperatingState = NodeOperatingState.OFF
@@ -1519,13 +1517,18 @@ class Node(SimComponent, ABC):
config: Node.ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema()) config: Node.ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema())
class ConfigSchema: class ConfigSchema(BaseModel, ABC):
"""Configuration Schema for Node based classes.""" """Configuration Schema for Node based classes."""
model_config = ConfigDict(arbitrary_types_allowed=True)
"""Configure pydantic to allow arbitrary types and to let the instance have attributes not present in the model."""
hostname: str
"The node hostname on the network."
revealed_to_red: bool = False revealed_to_red: bool = False
"Informs whether the node has been revealed to a red agent." "Informs whether the node has been revealed to a red agent."
start_up_duration: int = 3 start_up_duration: int = 0
"Time steps needed for the node to start up." "Time steps needed for the node to start up."
start_up_countdown: int = 0 start_up_countdown: int = 0
@@ -1549,8 +1552,9 @@ class Node(SimComponent, ABC):
red_scan_countdown: int = 0 red_scan_countdown: int = 0
"Time steps until reveal to red scan is complete." "Time steps until reveal to red scan is complete."
def from_config(cls, config: Dict) -> Node: @classmethod
"""Create Node object from a given configuration.""" def from_config(cls, config: Dict) -> "Node":
"""Create Node object from a given configuration dictionary."""
if config["type"] not in cls._registry: if config["type"] not in cls._registry:
msg = f"Configuration contains an invalid Node type: {config['type']}" msg = f"Configuration contains an invalid Node type: {config['type']}"
return ValueError(msg) return ValueError(msg)

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

@@ -42,6 +42,6 @@ class Computer(HostNode, identifier="computer"):
class ConfigSchema(HostNode.ConfigSchema): class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Computer class.""" """Configuration Schema for Computer class."""
pass hostname: str = "Computer"
pass pass

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

@@ -332,7 +332,7 @@ class HostNode(Node, identifier="HostNode"):
class ConfigSchema(Node.ConfigSchema): class ConfigSchema(Node.ConfigSchema):
"""Configuration Schema for HostNode class.""" """Configuration Schema for HostNode class."""
pass hostname: str = "HostNode"
def __init__(self, ip_address: IPV4Address, subnet_mask: IPV4Address, **kwargs): def __init__(self, ip_address: IPV4Address, subnet_mask: IPV4Address, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)

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

@@ -104,13 +104,14 @@ class Firewall(Router, identifier="firewall"):
class ConfigSchema(Router.ConfigSChema): class ConfigSchema(Router.ConfigSChema):
"""Configuration Schema for Firewall 'Nodes' within PrimAITE.""" """Configuration Schema for Firewall 'Nodes' within PrimAITE."""
pass hostname: str = "Firewall"
num_ports: int = 0
def __init__(self, hostname: str, **kwargs): def __init__(self, **kwargs):
if not kwargs.get("sys_log"): if not kwargs.get("sys_log"):
kwargs["sys_log"] = SysLog(hostname) kwargs["sys_log"] = SysLog(self.config.hostname)
super().__init__(hostname=hostname, num_ports=0, **kwargs) super().__init__(hostname=self.config.hostname, num_ports=self.config.num_ports, **kwargs)
self.connect_nic( self.connect_nic(
RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", port_name="external") RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", port_name="external")

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

@@ -1211,27 +1211,22 @@ class Router(NetworkNode, identifier="router"):
"The Router Interfaces on the node." "The Router Interfaces on the node."
network_interface: Dict[int, RouterInterface] = {} network_interface: Dict[int, RouterInterface] = {}
"The Router Interfaces on the node by port id." "The Router Interfaces on the node by port id."
acl: AccessControlList
route_table: RouteTable
config: "Router.ConfigSchema" = Field(default_factory=lambda: Router.ConfigSChema()) config: "Router.ConfigSchema" = Field(default_factory=lambda: Router.ConfigSchema())
class ConfigSChema(NetworkNode.ConfigSchema): class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Router Objects.""" """Configuration Schema for Router Objects."""
num_ports: int = 10 num_ports: int = 5
hostname: str = "Router" hostname: str = "Router"
ports: list = [] ports: list = []
sys_log: SysLog = SysLog(hostname)
acl: AccessControlList = AccessControlList(sys_log=sys_log, implicit_action=ACLAction.DENY, name=hostname)
route_table: RouteTable = RouteTable(sys_log=sys_log)
def __init__(self, hostname: str, num_ports: int = 5, **kwargs): def __init__(self, **kwargs):
if not kwargs.get("sys_log"): super().__init__(hostname=self.config.hostname, num_ports=self.config.num_ports, **kwargs)
kwargs["sys_log"] = SysLog(hostname) self.session_manager = RouterSessionManager(sys_log=self.config.sys_log)
if not kwargs.get("acl"):
kwargs["acl"] = AccessControlList(sys_log=kwargs["sys_log"], implicit_action=ACLAction.DENY, name=hostname)
if not kwargs.get("route_table"):
kwargs["route_table"] = RouteTable(sys_log=kwargs["sys_log"])
super().__init__(hostname=hostname, num_ports=num_ports, **kwargs)
self.session_manager = RouterSessionManager(sys_log=self.sys_log)
self.session_manager.node = self self.session_manager.node = self
self.software_manager.session_manager = self.session_manager self.software_manager.session_manager = self.session_manager
self.session_manager.software_manager = self.software_manager self.session_manager.software_manager = self.software_manager
@@ -1265,10 +1260,10 @@ class Router(NetworkNode, identifier="router"):
Initializes the router's ACL (Access Control List) with default rules, permitting essential protocols like ARP Initializes the router's ACL (Access Control List) with default rules, permitting essential protocols like ARP
and ICMP, which are necessary for basic network operations and diagnostics. and ICMP, which are necessary for basic network operations and diagnostics.
""" """
self.acl.add_rule( self.config.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22 action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
) )
self.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23) self.config.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
def setup_for_episode(self, episode: int): def setup_for_episode(self, episode: int):
""" """
@@ -1292,7 +1287,7 @@ class Router(NetworkNode, identifier="router"):
More information in user guide and docstring for SimComponent._init_request_manager. More information in user guide and docstring for SimComponent._init_request_manager.
""" """
rm = super()._init_request_manager() rm = super()._init_request_manager()
rm.add_request("acl", RequestType(func=self.acl._request_manager)) rm.add_request("acl", RequestType(func=self.config.acl._request_manager))
return rm return rm
def ip_is_router_interface(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool: def ip_is_router_interface(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool:
@@ -1346,7 +1341,7 @@ class Router(NetworkNode, identifier="router"):
""" """
state = super().describe_state() state = super().describe_state()
state["num_ports"] = self.config.num_ports state["num_ports"] = self.config.num_ports
state["acl"] = self.acl.describe_state() state["acl"] = self.config.acl.describe_state()
return state return state
def check_send_frame_to_session_manager(self, frame: Frame) -> bool: def check_send_frame_to_session_manager(self, frame: Frame) -> bool:
@@ -1393,7 +1388,7 @@ class Router(NetworkNode, identifier="router"):
return return
# Check if it's permitted # Check if it's permitted
permitted, rule = self.acl.is_permitted(frame) permitted, rule = self.config.acl.is_permitted(frame)
if not permitted: if not permitted:
at_port = self._get_port_of_nic(from_network_interface) at_port = self._get_port_of_nic(from_network_interface)
@@ -1566,83 +1561,18 @@ class Router(NetworkNode, identifier="router"):
) )
print(table) print(table)
# TODO: Remove - Cover normal config items with ConfigSchema. Move additional setup components to __init__ ? def setup_router(self, cfg: dict) -> Router:
""" TODO: This is the extra bit of Router's from_config metho. Needs sorting."""
@classmethod
def from_config(cls, cfg: dict, **kwargs) -> "Router":
"""Create a router based on a config dict.
Schema:
- hostname (str): unique name for this router.
- num_ports (int, optional): Number of network ports on the router. 8 by default
- ports (dict): Dict with integers from 1 - num_ports as keys. The values should be another dict specifying
ip_address and subnet_mask assigned to that ports (as strings)
- acl (dict): Dict with integers from 1 - max_acl_rules as keys. The key defines the position within the ACL
where the rule will be added (lower number is resolved first). The values should describe valid ACL
Rules as:
- action (str): either PERMIT or DENY
- src_port (str, optional): the named port such as HTTP, HTTPS, or POSTGRES_SERVER
- dst_port (str, optional): the named port such as HTTP, HTTPS, or POSTGRES_SERVER
- protocol (str, optional): the named IP protocol such as ICMP, TCP, or UDP
- src_ip_address (str, optional): IP address octet written in base 10
- dst_ip_address (str, optional): IP address octet written in base 10
- routes (list[dict]): List of route dicts with values:
- address (str): The destination address of the route.
- subnet_mask (str): The subnet mask of the route.
- next_hop_ip_address (str): The next hop IP for the route.
- metric (int): The metric of the route. Optional.
- default_route:
- next_hop_ip_address (str): The next hop IP for the route.
Example config:
```
{
'hostname': 'router_1',
'num_ports': 5,
'ports': {
1: {
'ip_address' : '192.168.1.1',
'subnet_mask' : '255.255.255.0',
},
2: {
'ip_address' : '192.168.0.1',
'subnet_mask' : '255.255.255.252',
}
},
'acl' : {
21: {'action': 'PERMIT', 'src_port': 'HTTP', dst_port: 'HTTP'},
22: {'action': 'PERMIT', 'src_port': 'ARP', 'dst_port': 'ARP'},
23: {'action': 'PERMIT', 'protocol': 'ICMP'},
},
'routes' : [
{'address': '192.168.0.0', 'subnet_mask': '255.255.255.0', 'next_hop_ip_address': '192.168.1.2'}
],
'default_route': {'next_hop_ip_address': '192.168.0.2'}
}
```
:param cfg: Router config adhering to schema described in main docstring body
:type cfg: dict
:return: Configured router.
:rtype: Router
"""
router = Router(
hostname=cfg["hostname"],
num_ports=int(cfg.get("num_ports", "5")),
operating_state=NodeOperatingState.ON
if not (p := cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
if "ports" in cfg: if "ports" in cfg:
for port_num, port_cfg in cfg["ports"].items(): for port_num, port_cfg in cfg["ports"].items():
router.configure_port( self.configure_port(
port=port_num, port=port_num,
ip_address=port_cfg["ip_address"], ip_address=port_cfg["ip_address"],
subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")), subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")),
) )
if "acl" in cfg: if "acl" in cfg:
for r_num, r_cfg in cfg["acl"].items(): for r_num, r_cfg in cfg["acl"].items():
router.acl.add_rule( self.config.acl.add_rule(
action=ACLAction[r_cfg["action"]], action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p], src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p], dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
@@ -1655,7 +1585,7 @@ class Router(NetworkNode, identifier="router"):
) )
if "routes" in cfg: if "routes" in cfg:
for route in cfg.get("routes"): for route in cfg.get("routes"):
router.route_table.add_route( self.config.route_table.add_route(
address=IPv4Address(route.get("address")), address=IPv4Address(route.get("address")),
subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")), subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")),
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")), next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
@@ -1664,5 +1594,5 @@ class Router(NetworkNode, identifier="router"):
if "default_route" in cfg: if "default_route" in cfg:
next_hop_ip_address = cfg["default_route"].get("next_hop_ip_address", None) next_hop_ip_address = cfg["default_route"].get("next_hop_ip_address", None)
if next_hop_ip_address: if next_hop_ip_address:
router.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address) self.config.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address)
return router return self

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

@@ -107,8 +107,9 @@ class Switch(NetworkNode, identifier="switch"):
class ConfigSchema(NetworkNode.ConfigSchema): class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Switch nodes within PrimAITE.""" """Configuration Schema for Switch nodes within PrimAITE."""
hostname: str = "Switch"
num_ports: int = 24 num_ports: int = 24
"The number of ports on the switch." "The number of ports on the switch. Default is 24."
def __init__(self, **kwargs): def __init__(self, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)

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

@@ -124,12 +124,12 @@ class WirelessRouter(Router, identifier="wireless_router"):
network_interface: Dict[int, Union[RouterInterface, WirelessAccessPoint]] = {} network_interface: Dict[int, Union[RouterInterface, WirelessAccessPoint]] = {}
airspace: AirSpace airspace: AirSpace
config: "WirelessRouter.ConfigSchema" = Field(default_factory=lambda: WirelessRouter.Configschema()) config: "WirelessRouter.ConfigSchema" = Field(default_factory=lambda: WirelessRouter.ConfigSchema())
class ConfigSchema(Router.ConfigSChema): class ConfigSchema(Router.ConfigSChema):
"""Configuration Schema for WirelessRouter nodes within PrimAITE.""" """Configuration Schema for WirelessRouter nodes within PrimAITE."""
pass hostname: str = "WirelessRouter"
def __init__(self, hostname: str, airspace: AirSpace, **kwargs): def __init__(self, hostname: str, airspace: AirSpace, **kwargs):
super().__init__(hostname=hostname, num_ports=0, airspace=airspace, **kwargs) super().__init__(hostname=hostname, num_ports=0, airspace=airspace, **kwargs)

3
tests/conftest.py
View File

@@ -195,12 +195,14 @@ def example_network() -> Network:
network = Network() network = Network()
# Router 1 # Router 1
# router_1 = Router(hostname="router_1", start_up_duration=0)
router_1 = Router(hostname="router_1", start_up_duration=0) router_1 = Router(hostname="router_1", start_up_duration=0)
router_1.power_on() router_1.power_on()
router_1.configure_port(port=1, ip_address="192.168.1.1", subnet_mask="255.255.255.0") router_1.configure_port(port=1, ip_address="192.168.1.1", subnet_mask="255.255.255.0")
router_1.configure_port(port=2, ip_address="192.168.10.1", subnet_mask="255.255.255.0") router_1.configure_port(port=2, ip_address="192.168.10.1", subnet_mask="255.255.255.0")
# Switch 1 # Switch 1
# switch_1_config = Switch.ConfigSchema()
switch_1 = Switch(hostname="switch_1", num_ports=8, start_up_duration=0) switch_1 = Switch(hostname="switch_1", num_ports=8, start_up_duration=0)
switch_1.power_on() switch_1.power_on()
@@ -208,6 +210,7 @@ def example_network() -> Network:
router_1.enable_port(1) router_1.enable_port(1)
# Switch 2 # Switch 2
# switch_2_config = Switch.ConfigSchema()
switch_2 = Switch(hostname="switch_2", num_ports=8, start_up_duration=0) switch_2 = Switch(hostname="switch_2", num_ports=8, start_up_duration=0)
switch_2.power_on() switch_2.power_on()
network.connect(endpoint_a=router_1.network_interface[2], endpoint_b=switch_2.network_interface[8]) network.connect(endpoint_a=router_1.network_interface[2], endpoint_b=switch_2.network_interface[8])