#1814: Remove hardcoded values + added test + remove unnecessary private parent attribute

This commit is contained in:
Czar Echavez
2023-09-06 11:12:03 +01:00
parent c349bb4484
commit d503e51c2d
6 changed files with 55 additions and 38 deletions

View File

@@ -1,6 +1,6 @@
"""Core of the PrimAITE Simulator.""" """Core of the PrimAITE Simulator."""
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Callable, Dict, List, Optional, Union from typing import Callable, Dict, List, Optional
from uuid import uuid4 from uuid import uuid4
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
@@ -140,7 +140,7 @@ class SimComponent(BaseModel):
kwargs["uuid"] = str(uuid4()) kwargs["uuid"] = str(uuid4())
super().__init__(**kwargs) super().__init__(**kwargs)
self._action_manager: ActionManager = self._init_action_manager() self._action_manager: ActionManager = self._init_action_manager()
self._parent: Optional["SimComponent"] = None self.parent: Optional["SimComponent"] = None
def _init_action_manager(self) -> ActionManager: def _init_action_manager(self) -> ActionManager:
""" """
@@ -213,24 +213,3 @@ class SimComponent(BaseModel):
Override this method with anything that needs to happen within the component for it to be reset. Override this method with anything that needs to happen within the component for it to be reset.
""" """
pass pass
@property
def parent(self) -> "SimComponent":
"""Reference to the parent object which manages this object.
:return: Parent object.
:rtype: SimComponent
"""
return self._parent
@parent.setter
def parent(self, new_parent: Union["SimComponent", None]) -> None:
if self._parent and new_parent:
msg = f"Overwriting parent of {self.uuid}. Old parent: {self._parent.uuid}, New parent: {new_parent.uuid}"
_LOGGER.warn(msg)
raise RuntimeWarning(msg)
self._parent = new_parent
@parent.deleter
def parent(self) -> None:
self._parent = None

View File

@@ -39,7 +39,7 @@ class SoftwareManager:
""" """
Add a Service to the manager. Add a Service to the manager.
:param service_class: The class of the service to add :param: service_class: The class of the service to add
""" """
service = service_class(software_manager=self, sys_log=self.sys_log) service = service_class(software_manager=self, sys_log=self.sys_log)

View File

@@ -1,8 +1,8 @@
from ipaddress import IPv4Address from ipaddress import IPv4Address
from typing import Any, Optional
from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port 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 from primaite.simulator.system.services.service import Service
@@ -20,9 +20,15 @@ class DataManipulatorService(Service):
kwargs["protocol"] = IPProtocol.TCP kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs) super().__init__(**kwargs)
def run(self): def start(self, target_ip_address: IPv4Address, payload: Optional[Any] = "DELETE TABLE users", **kwargs):
"""Run the DataManipulatorService actions.""" """
software_manager: SoftwareManager = self.software_manager Run the DataManipulatorService actions.
software_manager.send_payload_to_session_manager(
payload="SELECT * FROM users", dest_ip_address=IPv4Address("192.168.1.14"), dest_port=self.port :param: target_ip_address: The IP address of the target machine to attack
:param: payload: The payload to send to the target machine
"""
super().start()
self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=target_ip_address, dest_port=self.port
) )

View File

@@ -100,49 +100,49 @@ class Service(IOSoftware):
"""Stop the service.""" """Stop the service."""
_LOGGER.debug(f"Stopping service {self.name}") _LOGGER.debug(f"Stopping service {self.name}")
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.parent.sys_log.info(f"Stopping service {self.name}") self.sys_log.info(f"Stopping service {self.name}")
self.operating_state = ServiceOperatingState.STOPPED self.operating_state = ServiceOperatingState.STOPPED
def start(self) -> None: def start(self, **kwargs) -> None:
"""Start the service.""" """Start the service."""
_LOGGER.debug(f"Starting service {self.name}") _LOGGER.debug(f"Starting service {self.name}")
if self.operating_state == ServiceOperatingState.STOPPED: if self.operating_state == ServiceOperatingState.STOPPED:
self.parent.sys_log.info(f"Starting service {self.name}") self.sys_log.info(f"Starting service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING self.operating_state = ServiceOperatingState.RUNNING
def pause(self) -> None: def pause(self) -> None:
"""Pause the service.""" """Pause the service."""
_LOGGER.debug(f"Pausing service {self.name}") _LOGGER.debug(f"Pausing service {self.name}")
if self.operating_state == ServiceOperatingState.RUNNING: if self.operating_state == ServiceOperatingState.RUNNING:
self.parent.sys_log.info(f"Pausing service {self.name}") self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.PAUSED self.operating_state = ServiceOperatingState.PAUSED
def resume(self) -> None: def resume(self) -> None:
"""Resume paused service.""" """Resume paused service."""
_LOGGER.debug(f"Resuming service {self.name}") _LOGGER.debug(f"Resuming service {self.name}")
if self.operating_state == ServiceOperatingState.PAUSED: if self.operating_state == ServiceOperatingState.PAUSED:
self.parent.sys_log.info(f"Resuming service {self.name}") self.sys_log.info(f"Resuming service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING self.operating_state = ServiceOperatingState.RUNNING
def restart(self) -> None: def restart(self) -> None:
"""Restart running service.""" """Restart running service."""
_LOGGER.debug(f"Restarting service {self.name}") _LOGGER.debug(f"Restarting service {self.name}")
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
self.parent.sys_log.info(f"Pausing service {self.name}") self.sys_log.info(f"Pausing service {self.name}")
self.operating_state = ServiceOperatingState.RESTARTING self.operating_state = ServiceOperatingState.RESTARTING
self.restart_countdown = self.restarting_duration self.restart_countdown = self.restarting_duration
def disable(self) -> None: def disable(self) -> None:
"""Disable the service.""" """Disable the service."""
_LOGGER.debug(f"Disabling service {self.name}") _LOGGER.debug(f"Disabling service {self.name}")
self.parent.sys_log.info(f"Disabling Application {self.name}") self.sys_log.info(f"Disabling Application {self.name}")
self.operating_state = ServiceOperatingState.DISABLED self.operating_state = ServiceOperatingState.DISABLED
def enable(self) -> None: def enable(self) -> None:
"""Enable the disabled service.""" """Enable the disabled service."""
_LOGGER.debug(f"Enabling service {self.name}") _LOGGER.debug(f"Enabling service {self.name}")
if self.operating_state == ServiceOperatingState.DISABLED: if self.operating_state == ServiceOperatingState.DISABLED:
self.parent.sys_log.info(f"Enabling Application {self.name}") self.sys_log.info(f"Enabling Application {self.name}")
self.operating_state = ServiceOperatingState.STOPPED self.operating_state = ServiceOperatingState.STOPPED
def apply_timestep(self, timestep: int) -> None: def apply_timestep(self, timestep: int) -> None:

View File

@@ -0,0 +1,32 @@
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.networks import arcd_uc2_network
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.red_services.data_manipulator_service import DataManipulatorService
def test_creation():
network = arcd_uc2_network()
client_1: Node = network.get_node_by_hostname("client_1")
client_1.software_manager.add_service(service_class=DataManipulatorService)
data_manipulator_service: DataManipulatorService = client_1.software_manager.services["DataManipulatorBot"]
assert data_manipulator_service.name == "DataManipulatorBot"
assert data_manipulator_service.port == Port.POSTGRES_SERVER
assert data_manipulator_service.protocol == IPProtocol.TCP
# should have no session yet
assert len(client_1.session_manager.sessions_by_uuid) == 0
try:
data_manipulator_service.start(target_ip_address=IPv4Address("192.168.1.14"))
except Exception as e:
assert False, f"Test was not supposed to throw exception: {e}"
# there should be a session after the service is started
assert len(client_1.session_manager.sessions_by_uuid) == 1