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

#1814: initial implementation of data manipulator service

Browse Source
This commit is contained in:
Czar Echavez
2023-09-05 17:14:47 +01:00
parent 0493c2062c
commit c349bb4484
8 changed files with 163 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/network/transmission/transport_layer.py
View File

@@ -59,6 +59,8 @@ class Port(Enum):
"Alternative port for HTTP (HTTP_ALT) - Often used as an alternative HTTP port for web applications."
HTTPS_ALT = 8443
"Alternative port for HTTPS (HTTPS_ALT) - Used in some configurations for secure web traffic."
POSTGRES_SERVER = 5432
"Postgres SQL Server."
class UDPHeader(BaseModel):

95
src/primaite/simulator/system/core/session_manager.py
View File

@@ -1,12 +1,14 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.core import SimComponent
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.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader
if TYPE_CHECKING:
from primaite.simulator.network.hardware.base import ARPCache
@@ -135,7 +137,14 @@ class SessionManager:
dst_port = None
return protocol, src_ip_address, dst_ip_address, src_port, dst_port
def receive_payload_from_software_manager(self, payload: Any, session_id: Optional[int] = None):
def receive_payload_from_software_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
session_id: Optional[str] = None,
is_reattempt: bool = False,
) -> Union[Any, None]:
"""
Receive a payload from the SoftwareManager.
@@ -144,9 +153,50 @@ class SessionManager:
:param payload: The payload to be sent.
:param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
"""
# TODO: Implement session creation and
if session_id:
dest_ip_address = self.sessions_by_uuid[session_id].dst_ip_address
dest_port = self.sessions_by_uuid[session_id].dst_port
self.send_payload_to_nic(payload, session_id)
dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dest_ip_address)
if dst_mac_address:
outbound_nic = self.arp_cache.get_arp_cache_nic(dest_ip_address)
else:
if not is_reattempt:
self.arp_cache.send_arp_request(dest_ip_address)
return self.receive_payload_from_software_manager(
payload=payload,
dest_ip_address=dest_ip_address,
dest_port=dest_port,
session_id=session_id,
is_reattempt=True,
)
else:
return
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
ip=IPPacket(
src_ip_address=outbound_nic.ip_address,
dst_ip_address=dest_ip_address,
),
tcp=TCPHeader(
src_port=dest_port,
dst_port=dest_port,
),
payload=payload,
)
if not session_id:
session_key = self._get_session_key(frame, from_source=True)
session = self.sessions_by_key.get(session_key)
if not session:
# Create new session
session = Session.from_session_key(session_key)
self.sessions_by_key[session_key] = session
self.sessions_by_uuid[session.uuid] = session
outbound_nic.send_frame(frame)
def send_payload_to_software_manager(self, payload: Any, session_id: int):
"""
@@ -157,18 +207,6 @@ class SessionManager:
"""
self.software_manager.receive_payload_from_session_manger()
def send_payload_to_nic(self, payload: Any, session_id: int):
"""
Send a payload across the Network.
Takes a payload and a session_id. Builds a Frame and sends it across the network via a NIC.
:param payload: The payload to be sent.
:param session_id: The Session ID the payload originates from
"""
# TODO: Implement frame construction and sent to NIC.
pass
def receive_payload_from_nic(self, frame: Frame):
"""
Receive a Frame from the NIC.
@@ -187,3 +225,22 @@ class SessionManager:
self.sessions_by_uuid[session.uuid] = session
self.software_manager.receive_payload_from_session_manger(payload=frame, session=session)
# TODO: Implement the frame deconstruction and send to SoftwareManager.
def show(self, markdown: bool = False):
"""
Print tables describing the SessionManager.
Generate and print PrettyTable instances that show details about
session's destination IP Address, destination Ports and the protocol to use.
Output can be in Markdown format.
:param markdown: Use Markdown style in table output. Defaults to False.
"""
table = PrettyTable(["Destination IP", "Port", "Protocol"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Session Manager"
for session in self.sessions_by_key.values():
table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
print(table)

51
src/primaite/simulator/system/core/software_manager.py
View File

@@ -1,5 +1,8 @@
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
@@ -12,6 +15,10 @@ if TYPE_CHECKING:
from primaite.simulator.system.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from typing import Type, TypeVar
ServiceClass = TypeVar("ServiceClass", bound=Service)
class SoftwareManager:
"""A class that manages all running Services and Applications on a Node and facilitates their communication."""
@@ -28,18 +35,17 @@ class SoftwareManager:
self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {}
self.sys_log: SysLog = sys_log
def add_service(self, name: str, service: Service, port: Port, protocol: IPProtocol):
def add_service(self, service_class: Type[ServiceClass]):
"""
Add a Service to the manager.
:param name: The name of the service.
:param service: The service instance.
:param port: The port used by the service.
:param protocol: The network protocol used by the service.
:param service_class: The class of the service to add
"""
service = service_class(software_manager=self, sys_log=self.sys_log)
service.software_manager = self
self.services[name] = service
self.port_protocol_mapping[(port, protocol)] = service
self.services[service.name] = service
self.port_protocol_mapping[(service.port, service.protocol)] = service
def add_application(self, name: str, application: Application, port: Port, protocol: IPProtocol):
"""
@@ -75,14 +81,24 @@ class SoftwareManager:
else:
raise ValueError(f"No {target_software_type.name.lower()} found with the name {target_software}")
def send_payload_to_session_manger(self, payload: Any, session_id: Optional[int] = None):
def send_payload_to_session_manager(
self,
payload: Any,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
session_id: Optional[int] = None,
):
"""
Send a payload to the SessionManager.
:param payload: The payload to be sent.
:param dest_ip_address: The ip address of the payload destination.
:param dest_port: The port of the payload destination.
:param session_id: The Session ID the payload is to originate from. Optional.
"""
self.session_manager.receive_payload_from_software_manager(payload, session_id)
self.session_manager.receive_payload_from_software_manager(
payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
)
def receive_payload_from_session_manger(self, payload: Any, session: Session):
"""
@@ -97,3 +113,20 @@ class SoftwareManager:
# else:
# raise ValueError(f"No service or application found for port {port} and protocol {protocol}")
pass
def show(self, markdown: bool = False):
"""
Prints a table of the SwitchPorts on the Switch.
:param markdown: If True, outputs the table in markdown format. Default is False.
"""
table = PrettyTable(["Name", "Operating State", "Health State", "Port"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Software Manager"
for service in self.services.values():
table.add_row(
[service.name, service.operating_state.name, service.health_state_actual.name, service.port.value]
)
print(table)

28
src/primaite/simulator/system/services/red_services/data_manipulator_service.py Normal file
View File

@@ -0,0 +1,28 @@
from ipaddress import IPv4Address
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.service import Service
class DataManipulatorService(Service):
"""
Red Agent Data Integration Service.
The Service represents a bot that causes files/folders in the File System to
become corrupted.
"""
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulatorBot"
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def run(self):
"""Run the DataManipulatorService actions."""
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload="SELECT * FROM users", dest_ip_address=IPv4Address("192.168.1.14"), dest_port=self.port
)

4
src/primaite/simulator/system/services/service.py
View File

@@ -1,4 +1,3 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, Optional
@@ -33,7 +32,7 @@ class Service(IOSoftware):
Services are programs that run in the background and may perform input/output operations.
"""
operating_state: ServiceOperatingState
operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED
"The current operating state of the Service."
restart_duration: int = 5
"How many timesteps does it take to restart this service."
@@ -51,7 +50,6 @@ class Service(IOSoftware):
am.add_action("enable", Action(func=lambda request, context: self.enable()))
return am
@abstractmethod
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

18
src/primaite/simulator/system/software.py
View File

@@ -1,9 +1,10 @@
from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, Set
from typing import Any, Dict
from primaite.simulator.core import Action, ActionManager, SimComponent
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.sys_log import SysLog
class SoftwareType(Enum):
@@ -62,11 +63,11 @@ class Software(SimComponent):
name: str
"The name of the software."
health_state_actual: SoftwareHealthState
health_state_actual: SoftwareHealthState = SoftwareHealthState.GOOD
"The actual health state of the software."
health_state_visible: SoftwareHealthState
health_state_visible: SoftwareHealthState = SoftwareHealthState.GOOD
"The health state of the software visible to the red agent."
criticality: SoftwareCriticality
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
"The criticality level of the software."
patching_count: int = 0
"The count of patches applied to the software, defaults to 0."
@@ -74,6 +75,10 @@ class Software(SimComponent):
"The count of times the software has been scanned, defaults to 0."
revealed_to_red: bool = False
"Indicates if the software has been revealed to red agent, defaults is False."
software_manager: Any = None
"An instance of Software Manager that is used by the parent node."
sys_log: SysLog = None
"An instance of SysLog that is used by the parent node."
def _init_action_manager(self) -> ActionManager:
am = super()._init_action_manager()
@@ -132,7 +137,6 @@ class Software(SimComponent):
"""
self.health_state_actual = health_state
@abstractmethod
def install(self) -> None:
"""
Perform first-time setup of this service on a node.
@@ -175,8 +179,8 @@ class IOSoftware(Software):
"Indicates if the software uses TCP protocol for communication. Default is True."
udp: bool = True
"Indicates if the software uses UDP protocol for communication. Default is True."
ports: Set[Port]
"The set of ports to which the software is connected."
port: Port
"The port to which the software is connected."
@abstractmethod
def describe_state(self) -> Dict:

8
tests/integration_tests/system/test_database_on_node.py
View File

@@ -11,9 +11,7 @@ def test_installing_database():
health_state_actual=SoftwareHealthState.GOOD,
health_state_visible=SoftwareHealthState.GOOD,
criticality=SoftwareCriticality.MEDIUM,
ports=[
Port.SQL_SERVER,
],
port=Port.SQL_SERVER,
operating_state=ServiceOperatingState.RUNNING,
)
@@ -40,9 +38,7 @@ def test_uninstalling_database():
health_state_actual=SoftwareHealthState.GOOD,
health_state_visible=SoftwareHealthState.GOOD,
criticality=SoftwareCriticality.MEDIUM,
ports=[
Port.SQL_SERVER,
],
port=Port.SQL_SERVER,
operating_state=ServiceOperatingState.RUNNING,
)

4
tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py
View File

@@ -10,8 +10,6 @@ def test_creation():
health_state_actual=SoftwareHealthState.GOOD,
health_state_visible=SoftwareHealthState.GOOD,
criticality=SoftwareCriticality.MEDIUM,
ports=[
Port.SQL_SERVER,
],
port=Port.SQL_SERVER,
operating_state=ServiceOperatingState.RUNNING,
)