#1814: Remove hardcoded values + added test + remove unnecessary private parent attribute
This commit is contained in:
@@ -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
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user