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

Merge remote-tracking branch 'origin/dev' into feature/2646_Update-pre-commit-to-check-for-valid-copyright

Browse Source
This commit is contained in:
Marek Wolan
2024-06-13 12:52:09 +01:00
parent 6762fd805e ae2bbd3e00
commit fad8d936e6
54 changed files with 9740 additions and 647 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,6 +29,7 @@ from primaite.simulator.network.nmne import (
NMNE_CAPTURE_KEYWORDS,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.packet_capture import PacketCapture
from primaite.simulator.system.core.session_manager import SessionManager
@@ -37,6 +38,7 @@ from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.processes.process import Process
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import IOSoftware
from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
from primaite.utils.validators import IPV4Address
IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
@@ -108,10 +110,14 @@ class NetworkInterface(SimComponent, ABC):
nmne: Dict = Field(default_factory=lambda: {})
"A dict containing details of the number of malicious network events captured."
traffic: Dict = Field(default_factory=lambda: {})
"A dict containing details of the inbound and outbound traffic by port and protocol."
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().setup_for_episode(episode=episode)
self.nmne = {}
self.traffic = {}
if episode and self.pcap and SIM_OUTPUT.save_pcap_logs:
self.pcap.current_episode = episode
self.pcap.setup_logger()
@@ -147,6 +153,7 @@ class NetworkInterface(SimComponent, ABC):
)
if CAPTURE_NMNE:
state.update({"nmne": {k: v for k, v in self.nmne.items()}})
state.update({"traffic": convert_dict_enum_keys_to_enum_values(self.traffic)})
return state
@abstractmethod
@@ -237,6 +244,47 @@ class NetworkInterface(SimComponent, ABC):
# Increment a generic counter if keyword capturing is not enabled
keyword_level["*"] = keyword_level.get("*", 0) + 1
def _capture_traffic(self, frame: Frame, inbound: bool = True):
"""
Capture traffic statistics at the Network Interface.
:param frame: The network frame containing the traffic data.
:type frame: Frame
:param inbound: Flag indicating if the traffic is inbound or outbound. Defaults to True.
:type inbound: bool
"""
# Determine the direction of the traffic
direction = "inbound" if inbound else "outbound"
# Initialize protocol and port variables
protocol = None
port = None
# Identify the protocol and port from the frame
if frame.tcp:
protocol = IPProtocol.TCP
port = frame.tcp.dst_port
elif frame.udp:
protocol = IPProtocol.UDP
port = frame.udp.dst_port
elif frame.icmp:
protocol = IPProtocol.ICMP
# Ensure the protocol is in the capture dict
if protocol not in self.traffic:
self.traffic[protocol] = {}
# Handle non-ICMP protocols that use ports
if protocol != IPProtocol.ICMP:
if port not in self.traffic[protocol]:
self.traffic[protocol][port] = {"inbound": 0, "outbound": 0}
self.traffic[protocol][port][direction] += frame.size
else:
# Handle ICMP protocol separately (ICMP does not use ports)
if not self.traffic[protocol]:
self.traffic[protocol] = {"inbound": 0, "outbound": 0}
self.traffic[protocol][direction] += frame.size
@abstractmethod
def send_frame(self, frame: Frame) -> bool:
"""
@@ -246,6 +294,7 @@ class NetworkInterface(SimComponent, ABC):
:return: A boolean indicating whether the frame was successfully sent.
"""
self._capture_nmne(frame, inbound=False)
self._capture_traffic(frame, inbound=False)
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
@@ -256,6 +305,7 @@ class NetworkInterface(SimComponent, ABC):
:return: A boolean indicating whether the frame was successfully received.
"""
self._capture_nmne(frame, inbound=True)
self._capture_traffic(frame, inbound=True)
def __str__(self) -> str:
"""
@@ -767,6 +817,24 @@ class Node(SimComponent):
self.session_manager.software_manager = self.software_manager
self._install_system_software()
def ip_is_network_interface(self, ip_address: IPv4Address, enabled_only: bool = False) -> bool:
"""
Checks if a given IP address belongs to any of the nodes interfaces.
:param ip_address: The IP address to check.
:param enabled_only: If True, only considers enabled network interfaces.
:return: True if the IP address is assigned to one of the nodes interfaces; False otherwise.
"""
for network_interface in self.network_interface.values():
if not hasattr(network_interface, "ip_address"):
continue
if network_interface.ip_address == ip_address:
if enabled_only:
return network_interface.enabled
else:
return True
return False
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().setup_for_episode(episode=episode)

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

@@ -8,6 +8,8 @@ from primaite import getLogger
from primaite.simulator.network.hardware.base import IPWiredNetworkInterface, Link, Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.nmap import NMAP
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.arp.arp import ARP, ARPPacket
from primaite.simulator.system.services.dns.dns_client import DNSClient
@@ -303,6 +305,7 @@ class HostNode(Node):
"DNSClient": DNSClient,
"NTPClient": NTPClient,
"WebBrowser": WebBrowser,
"NMAP": NMAP,
}
"""List of system software that is automatically installed on nodes."""
@@ -315,6 +318,16 @@ class HostNode(Node):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
@property
def nmap(self) -> Optional[NMAP]:
"""
Return the NMAP application installed on the Node.
:return: NMAP application installed on the Node.
:rtype: Optional[NMAP]
"""
return self.software_manager.software.get("NMAP")
@property
def arp(self) -> Optional[ARP]:
"""
@@ -366,8 +379,15 @@ class HostNode(Node):
elif frame.udp:
dst_port = frame.udp.dst_port
can_accept_nmap = False
if self.software_manager.software.get("NMAP"):
if self.software_manager.software["NMAP"].operating_state == ApplicationOperatingState.RUNNING:
can_accept_nmap = True
accept_nmap = can_accept_nmap and frame.payload.__class__.__name__ == "PortScanPayload"
accept_frame = False
if frame.icmp or dst_port in self.software_manager.get_open_ports():
if frame.icmp or dst_port in self.software_manager.get_open_ports() or accept_nmap:
# accept the frame as the port is open or if it's an ICMP frame
accept_frame = True

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

@@ -19,6 +19,7 @@ from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.nmap import NMAP
from primaite.simulator.system.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.services.arp.arp import ARP
@@ -1239,6 +1240,7 @@ class Router(NetworkNode):
icmp.router = self
self.software_manager.install(RouterARP)
self.arp.router = self
self.software_manager.install(NMAP)
def _set_default_acl(self):
"""