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

Make all requests return a RequestResponse

Browse Source
This commit is contained in:
Marek Wolan
2024-03-08 14:58:34 +00:00
parent d331224b45
commit beb51834f9
13 changed files with 275 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -12,6 +12,7 @@ from pydantic import BaseModel, Field
from primaite import getLogger
from primaite.exceptions import NetworkError
from primaite.interface.request import RequestResponse
from primaite.simulator import SIM_OUTPUT
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.domain.account import Account
@@ -115,8 +116,8 @@ class NetworkInterface(SimComponent, ABC):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request("enable", RequestType(func=lambda request, context: self.enable()))
rm.add_request("disable", RequestType(func=lambda request, context: self.disable()))
rm.add_request("enable", RequestType(func=lambda request, context: RequestResponse.from_bool(self.enable())))
rm.add_request("disable", RequestType(func=lambda request, context: RequestResponse.from_bool(self.disable())))
return rm
@@ -140,14 +141,16 @@ class NetworkInterface(SimComponent, ABC):
return state
@abstractmethod
def enable(self):
def enable(self) -> bool:
"""Enable the interface."""
pass
return False
@abstractmethod
def disable(self):
def disable(self) -> bool:
"""Disable the interface."""
pass
return False
def _capture_nmne(self, frame: Frame, inbound: bool = True) -> None:
"""
@@ -783,16 +786,28 @@ class Node(SimComponent):
self._application_request_manager = RequestManager()
rm.add_request("application", RequestType(func=self._application_request_manager))
rm.add_request("scan", RequestType(func=lambda request, context: self.reveal_to_red()))
rm.add_request(
"scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.reveal_to_red()))
)
rm.add_request("shutdown", RequestType(func=lambda request, context: self.power_off()))
rm.add_request("startup", RequestType(func=lambda request, context: self.power_on()))
rm.add_request("reset", RequestType(func=lambda request, context: self.reset())) # TODO implement node reset
rm.add_request("logon", RequestType(func=lambda request, context: ...)) # TODO implement logon request
rm.add_request("logoff", RequestType(func=lambda request, context: ...)) # TODO implement logoff request
rm.add_request(
"shutdown", RequestType(func=lambda request, context: RequestResponse.from_bool(self.power_off()))
)
rm.add_request("startup", RequestType(func=lambda request, context: RequestResponse.from_bool(self.power_on())))
rm.add_request(
"reset", RequestType(func=lambda request, context: RequestResponse.from_bool(self.reset()))
) # TODO implement node reset
rm.add_request(
"logon", RequestType(func=lambda request, context: RequestResponse.from_bool(False))
) # TODO implement logon request
rm.add_request(
"logoff", RequestType(func=lambda request, context: RequestResponse.from_bool(False))
) # TODO implement logoff request
self._os_request_manager = RequestManager()
self._os_request_manager.add_request("scan", RequestType(func=lambda request, context: self.scan()))
self._os_request_manager.add_request(
"scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan()))
)
rm.add_request("os", RequestType(func=self._os_request_manager))
return rm
@@ -973,7 +988,7 @@ class Node(SimComponent):
self.file_system.apply_timestep(timestep=timestep)
def scan(self) -> None:
def scan(self) -> bool:
"""
Scan the node and all the items within it.
@@ -987,8 +1002,9 @@ class Node(SimComponent):
to the red agent.
"""
self.node_scan_countdown = self.node_scan_duration
return True
def reveal_to_red(self) -> None:
def reveal_to_red(self) -> bool:
"""
Reveals the node and all the items within it to the red agent.
@@ -1002,34 +1018,40 @@ class Node(SimComponent):
`revealed_to_red` to `True`.
"""
self.red_scan_countdown = self.node_scan_duration
return True
def power_on(self):
def power_on(self) -> bool:
"""Power on the Node, enabling its NICs if it is in the OFF state."""
if self.operating_state == NodeOperatingState.OFF:
self.operating_state = NodeOperatingState.BOOTING
self.start_up_countdown = self.start_up_duration
if self.start_up_duration <= 0:
self.operating_state = NodeOperatingState.ON
self._start_up_actions()
self.sys_log.info("Power on")
for network_interface in self.network_interfaces.values():
network_interface.enable()
return True
if self.operating_state == NodeOperatingState.OFF:
self.operating_state = NodeOperatingState.BOOTING
self.start_up_countdown = self.start_up_duration
return True
def power_off(self):
return False
def power_off(self) -> bool:
"""Power off the Node, disabling its NICs if it is in the ON state."""
if self.shut_down_duration <= 0:
self._shut_down_actions()
self.operating_state = NodeOperatingState.OFF
self.sys_log.info("Power off")
return True
if self.operating_state == NodeOperatingState.ON:
for network_interface in self.network_interfaces.values():
network_interface.disable()
self.operating_state = NodeOperatingState.SHUTTING_DOWN
self.shut_down_countdown = self.shut_down_duration
return True
return False
if self.shut_down_duration <= 0:
self._shut_down_actions()
self.operating_state = NodeOperatingState.OFF
self.sys_log.info("Power off")
def reset(self):
def reset(self) -> bool:
"""
Resets the node.
@@ -1040,6 +1062,8 @@ class Node(SimComponent):
self.is_resetting = True
self.sys_log.info("Resetting")
self.power_off()
return True
return False
def connect_nic(self, network_interface: NetworkInterface, port_name: Optional[str] = None):
"""

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

@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.network.hardware.base import IPWiredNetworkInterface
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
@@ -308,19 +309,24 @@ class AccessControlList(SimComponent):
rm.add_request(
"add_rule",
RequestType(
func=lambda request, context: self.add_rule(
action=ACLAction[request[0]],
protocol=None if request[1] == "ALL" else IPProtocol[request[1]],
src_ip_address=None if request[2] == "ALL" else IPv4Address(request[2]),
src_port=None if request[3] == "ALL" else Port[request[3]],
dst_ip_address=None if request[4] == "ALL" else IPv4Address(request[4]),
dst_port=None if request[5] == "ALL" else Port[request[5]],
position=int(request[6]),
func=lambda request, context: RequestResponse.from_bool(
self.add_rule(
action=ACLAction[request[0]],
protocol=None if request[1] == "ALL" else IPProtocol[request[1]],
src_ip_address=None if request[2] == "ALL" else IPv4Address(request[2]),
src_port=None if request[3] == "ALL" else Port[request[3]],
dst_ip_address=None if request[4] == "ALL" else IPv4Address(request[4]),
dst_port=None if request[5] == "ALL" else Port[request[5]],
position=int(request[6]),
)
)
),
)
rm.add_request("remove_rule", RequestType(func=lambda request, context: self.remove_rule(int(request[0]))))
rm.add_request(
"remove_rule",
RequestType(func=lambda request, context: RequestResponse.from_bool(self.remove_rule(int(request[0])))),
)
return rm
def describe_state(self) -> Dict:
@@ -366,7 +372,7 @@ class AccessControlList(SimComponent):
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
position: int = 0,
) -> None:
) -> bool:
"""
Adds a new ACL rule to control network traffic based on specified criteria.
@@ -423,10 +429,12 @@ class AccessControlList(SimComponent):
src_port=src_port,
dst_port=dst_port,
)
return True
else:
raise ValueError(f"Cannot add ACL rule, position {position} is out of bounds.")
return False
def remove_rule(self, position: int) -> None:
def remove_rule(self, position: int) -> bool:
"""
Remove an ACL rule from a specific position.
@@ -437,8 +445,10 @@ class AccessControlList(SimComponent):
rule = self._acl[position] # noqa
self._acl[position] = None
del rule
return True
else:
raise ValueError(f"Cannot remove ACL rule, position {position} is out of bounds.")
return False
def is_permitted(self, frame: Frame) -> Tuple[bool, ACLRule]:
"""Check if a packet with the given properties is permitted through the ACL."""