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

Merge remote-tracking branch 'origin/dev' into feature/2257-router-routes-cannot-be-represented-in-config-file

Browse Source
This commit is contained in:
Czar Echavez
2024-02-29 11:34:56 +00:00
parent 9a4587155b 312b5c35bd
commit eefc2739c8
58 changed files with 548 additions and 963 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/primaite/simulator/network/airspace.py
View File

@@ -273,11 +273,6 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def enable(self):
"""
Enables this wired network interface and attempts to send a "hello" message to the default gateway.

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

@@ -45,19 +45,12 @@ class Network(SimComponent):
self._nx_graph = MultiGraph()
def set_original_state(self):
"""Sets the original state."""
for node in self.nodes.values():
node.set_original_state()
for link in self.links.values():
link.set_original_state()
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
for node in self.nodes.values():
node.reset_component_for_episode(episode)
node.setup_for_episode(episode=episode)
for link in self.links.values():
link.reset_component_for_episode(episode)
link.setup_for_episode(episode=episode)
for node in self.nodes.values():
node.power_on()
@@ -179,7 +172,7 @@ class Network(SimComponent):
def clear_links(self):
"""Clear all the links in the network by resetting their component state for the episode."""
for link in self.links.values():
link.reset_component_for_episode()
link.setup_for_episode(episode=0) # TODO: shouldn't be using this method here.
def draw(self, seed: int = 123):
"""

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

@@ -100,6 +100,15 @@ class NetworkInterface(SimComponent, ABC):
nmne: Dict = Field(default_factory=lambda: {})
"A dict containing details of the number of malicious network events captured."
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().setup_for_episode(episode=episode)
self.nmne = {}
if episode and self.pcap:
self.pcap.current_episode = episode
self.pcap.setup_logger()
self.enable()
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
@@ -127,15 +136,6 @@ class NetworkInterface(SimComponent, ABC):
state.update({"nmne": self.nmne})
return state
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().reset_component_for_episode(episode)
self.nmne = {}
if episode and self.pcap:
self.pcap.current_episode = episode
self.pcap.setup_logger()
self.enable()
@abstractmethod
def enable(self):
"""Enable the interface."""
@@ -547,14 +547,6 @@ class Link(SimComponent):
self.endpoint_b.connect_link(self)
self.endpoint_up()
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"bandwidth", "current_load"}
self._original_state = self.model_dump(include=vals_to_include)
super().set_original_state()
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -740,50 +732,20 @@ class Node(SimComponent):
self.session_manager.node = self
self.session_manager.software_manager = self.software_manager
self._install_system_software()
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
for software in self.software_manager.software.values():
software.set_original_state()
self.file_system.set_original_state()
for network_interface in self.network_interfaces.values():
network_interface.set_original_state()
vals_to_include = {
"hostname",
"default_gateway",
"operating_state",
"revealed_to_red",
"start_up_duration",
"start_up_countdown",
"shut_down_duration",
"shut_down_countdown",
"is_resetting",
"node_scan_duration",
"node_scan_countdown",
"red_scan_countdown",
}
self._original_state = self.model_dump(include=vals_to_include)
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().reset_component_for_episode(episode)
# Reset Session Manager
self.session_manager.clear()
super().setup_for_episode(episode=episode)
# Reset File System
self.file_system.reset_component_for_episode(episode)
self.file_system.setup_for_episode(episode=episode)
# Reset all Nics
for network_interface in self.network_interfaces.values():
network_interface.reset_component_for_episode(episode)
network_interface.setup_for_episode(episode=episode)
for software in self.software_manager.software.values():
software.reset_component_for_episode(episode)
software.setup_for_episode(episode=episode)
if episode and self.sys_log:
self.sys_log.current_episode = episode

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

@@ -209,11 +209,6 @@ class NIC(IPWiredNetworkInterface):
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def receive_frame(self, frame: Frame) -> bool:
"""
Attempt to receive and process a network frame from the connected Link.

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

@@ -111,24 +111,6 @@ class Firewall(Router):
sys_log=kwargs["sys_log"], implicit_action=ACLAction.PERMIT, name=f"{hostname} - External Outbound"
)
self.set_original_state()
def set_original_state(self):
"""Set the original state for the Firewall."""
super().set_original_state()
vals_to_include = {
"internal_port",
"external_port",
"dmz_port",
"internal_inbound_acl",
"internal_outbound_acl",
"dmz_inbound_acl",
"dmz_outbound_acl",
"external_inbound_acl",
"external_outbound_acl",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def describe_state(self) -> Dict:
"""
Describes the current state of the Firewall.

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

@@ -136,11 +136,6 @@ class ACLRule(SimComponent):
rule_strings.append(f"{key}={value}")
return ", ".join(rule_strings)
def set_original_state(self):
"""Sets the original state."""
vals_to_keep = {"action", "protocol", "src_ip_address", "src_port", "dst_ip_address", "dst_port"}
self._original_state = self.model_dump(include=vals_to_keep, exclude_none=True)
def describe_state(self) -> Dict:
"""
Describes the current state of the ACLRule.
@@ -296,48 +291,6 @@ class AccessControlList(SimComponent):
super().__init__(**kwargs)
self._acl = [None] * (self.max_acl_rules - 1)
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
self.implicit_rule.set_original_state()
vals_to_keep = {"implicit_action", "max_acl_rules", "acl"}
self._original_state = self.model_dump(include=vals_to_keep, exclude_none=True)
for i, rule in enumerate(self._acl):
if not rule:
continue
self._default_config[i] = {"action": rule.action.name}
if rule.src_ip_address:
self._default_config[i]["src_ip"] = str(rule.src_ip_address)
if rule.dst_ip_address:
self._default_config[i]["dst_ip"] = str(rule.dst_ip_address)
if rule.src_port:
self._default_config[i]["src_port"] = rule.src_port.name
if rule.dst_port:
self._default_config[i]["dst_port"] = rule.dst_port.name
if rule.protocol:
self._default_config[i]["protocol"] = rule.protocol.name
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.implicit_rule.reset_component_for_episode(episode)
super().reset_component_for_episode(episode)
self._reset_rules_to_default()
def _reset_rules_to_default(self) -> None:
"""Clear all ACL rules and set them to the default rules config."""
self._acl = [None] * (self.max_acl_rules - 1)
for r_num, r_cfg in self._default_config.items():
self.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
position=r_num,
)
def _init_request_manager(self) -> RequestManager:
# TODO: Add src and dst wildcard masks as positional args in this request.
@@ -616,11 +569,6 @@ class RouteEntry(SimComponent):
metric: float = 0.0
"The cost metric for this route. Default is 0.0."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"address", "subnet_mask", "next_hop_ip_address", "metric"}
self._original_values = self.model_dump(include=vals_to_include)
def describe_state(self) -> Dict:
"""
Describes the current state of the RouteEntry.
@@ -653,17 +601,6 @@ class RouteTable(SimComponent):
default_route: Optional[RouteEntry] = None
sys_log: SysLog
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
self._original_state["routes_orig"] = self.routes
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.routes.clear()
self.routes = self._original_state["routes_orig"]
super().reset_component_for_episode(episode)
def describe_state(self) -> Dict:
"""
Describes the current state of the RouteTable.
@@ -1104,8 +1041,6 @@ class Router(NetworkNode):
self._set_default_acl()
self.set_original_state()
def _install_system_software(self):
"""
Installs essential system software and network services on the router.
@@ -1131,20 +1066,7 @@ class Router(NetworkNode):
self.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
self.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
def set_original_state(self):
"""
Sets or resets the router to its original configuration state.
This method is called to initialize the router's state or to revert it to a known good configuration during
network simulations or after configuration changes.
"""
self.acl.set_original_state()
self.route_table.set_original_state()
super().set_original_state()
vals_to_include = {"num_ports"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""
Resets the router's components for a new network simulation episode.
@@ -1154,13 +1076,10 @@ class Router(NetworkNode):
:param episode: The episode number for which the router is being reset.
"""
self.software_manager.arp.clear()
self.acl.reset_component_for_episode(episode)
self.route_table.reset_component_for_episode(episode)
for i, network_interface in self.network_interface.items():
network_interface.reset_component_for_episode(episode)
for i, _ in self.network_interface.items():
self.enable_port(i)
super().reset_component_for_episode(episode)
super().setup_for_episode(episode=episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
@@ -1391,7 +1310,6 @@ class Router(NetworkNode):
network_interface.ip_address = ip_address
network_interface.subnet_mask = subnet_mask
self.sys_log.info(f"Configured Network Interface {network_interface}")
self.set_original_state()
def enable_port(self, port: int):
"""
@@ -1493,8 +1411,16 @@ class Router(NetworkNode):
subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")),
)
if "acl" in cfg:
router.acl._default_config = cfg["acl"] # save the config to allow resetting
router.acl._reset_rules_to_default() # read the config and apply rules
for r_num, r_cfg in cfg["acl"].items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
position=r_num,
)
if "routes" in cfg:
for route in cfg.get("routes"):
router.route_table.add_route(

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

@@ -32,12 +32,6 @@ class SwitchPort(WiredNetworkInterface):
_connected_node: Optional[Switch] = None
"The Switch to which the SwitchPort is connected."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"port_num", "mac_address", "speed", "mtu", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
super().set_original_state()
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

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

@@ -122,8 +122,6 @@ class WirelessRouter(Router):
self.connect_nic(RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
self.set_original_state()
@property
def wireless_access_point(self) -> WirelessAccessPoint:
"""
@@ -166,7 +164,6 @@ class WirelessRouter(Router):
network_interface.ip_address = ip_address
network_interface.subnet_mask = subnet_mask
self.sys_log.info(f"Configured WAP {network_interface}")
self.set_original_state()
self.wireless_access_point.frequency = frequency # Set operating frequency
self.wireless_access_point.enable() # Re-enable the WAP with new settings