Merge branch 'dev' into feature/2041_2042-Add-NTP-Services

src/primaite/simulator/system/services/database/database_service.py

@@ -2,6 +2,7 @@ from datetime import datetime
 from ipaddress import IPv4Address
 from typing import Any, Dict, List, Literal, Optional, Union
 
+from primaite import getLogger
 from primaite.simulator.file_system.file_system import File
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.network.transmission.transport_layer import Port
@@ -10,6 +11,8 @@ from primaite.simulator.system.services.ftp.ftp_client import FTPClient
 from primaite.simulator.system.services.service import Service, ServiceOperatingState
 from primaite.simulator.system.software import SoftwareHealthState
 
+_LOGGER = getLogger(__name__)
+
 
 class DatabaseService(Service):
     """
@@ -38,6 +41,25 @@ class DatabaseService(Service):
         self._db_file: File
         self._create_db_file()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting DatabaseService original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {
+            "password",
+            "connections",
+            "backup_server",
+            "latest_backup_directory",
+            "latest_backup_file_name",
+        }
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug("Resetting DatabaseService original state on node {self.software_manager.node.hostname}")
+        self.connections.clear()
+        super().reset_component_for_episode(episode)
+
     def configure_backup(self, backup_server: IPv4Address):
         """
         Set up the database backup.
@@ -48,13 +70,17 @@ class DatabaseService(Service):
 
     def backup_database(self) -> bool:
         """Create a backup of the database to the configured backup server."""
+        # check if this action can be performed
+        if not self._can_perform_action():
+            return False
+
         # check if the backup server was configured
         if self.backup_server is None:
             self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.")
             return False
 
         software_manager: SoftwareManager = self.software_manager
-        ftp_client_service: FTPClient = software_manager.software["FTPClient"]
+        ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
 
         # send backup copy of database file to FTP server
         response = ftp_client_service.send_file(
@@ -73,8 +99,12 @@ class DatabaseService(Service):
 
     def restore_backup(self) -> bool:
         """Restore a backup from backup server."""
+        # check if this action can be performed
+        if not self._can_perform_action():
+            return False
+
         software_manager: SoftwareManager = self.software_manager
-        ftp_client_service: FTPClient = software_manager.software["FTPClient"]
+        ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
 
         # retrieve backup file from backup server
         response = ftp_client_service.request_file(
@@ -173,6 +203,9 @@ class DatabaseService(Service):
         :param session_id: The session identifier.
         :return: True if the Status Code is 200, otherwise False.
         """
+        if not super().receive(payload=payload, session_id=session_id, **kwargs):
+            return False
+
         result = {"status_code": 500, "data": []}
         if isinstance(payload, dict) and payload.get("type"):
             if payload["type"] == "connect_request":

src/primaite/simulator/system/services/dns/dns_client.py

@@ -29,6 +29,18 @@ class DNSClient(Service):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting DNSClient original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"dns_server"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        self.dns_cache.clear()
+        super().reset_component_for_episode(episode)
+
     def describe_state(self) -> Dict:
         """
         Describes the current state of the software.
@@ -42,23 +54,18 @@ class DNSClient(Service):
         state = super().describe_state()
         return state
 
-    def reset_component_for_episode(self, episode: int):
-        """
-        Resets the Service component for a new episode.
-
-        This method ensures the Service is ready for a new episode, including resetting any
-        stateful properties or statistics, and clearing any message queues.
-        """
-        pass
-
-    def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address):
+    def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address) -> bool:
         """
         Adds a domain name to the DNS Client cache.
 
         :param: domain_name: The domain name to save to cache
         :param: ip_address: The IP Address to attach the domain name to
         """
+        if not self._can_perform_action():
+            return False
+
         self.dns_cache[domain_name] = ip_address
+        return True
 
     def check_domain_exists(
         self,
@@ -72,6 +79,9 @@ class DNSClient(Service):
         :param: session_id: The Session ID the payload is to originate from. Optional.
         :param: is_reattempt: Checks if the request has been reattempted. Default is False.
         """
+        if not self._can_perform_action():
+            return False
+
         # check if DNS server is configured
         if self.dns_server is None:
             self.sys_log.error(f"{self.name}: DNS Server is not configured")

src/primaite/simulator/system/services/dns/dns_server.py

@@ -28,6 +28,20 @@ class DNSServer(Service):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting DNSServer original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"dns_table"}
+        self._original_state["dns_table_orig"] = self.model_dump(include=vals_to_include)["dns_table"]
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        self.dns_table.clear()
+        for key, value in self._original_state["dns_table_orig"].items():
+            self.dns_table[key] = value
+        super().reset_component_for_episode(episode)
+
     def describe_state(self) -> Dict:
         """
         Describes the current state of the software.
@@ -48,6 +62,9 @@ class DNSServer(Service):
         :param target_domain: The single domain name requested by a DNS client.
         :return ip_address: The IP address of that domain name or None.
         """
+        if not self._can_perform_action():
+            return
+
         return self.dns_table.get(target_domain)
 
     def dns_register(self, domain_name: str, domain_ip_address: IPv4Address):
@@ -60,17 +77,11 @@ class DNSServer(Service):
         :param: domain_ip_address: The IP address that the domain should route to
         :type: domain_ip_address: IPv4Address
         """
+        if not self._can_perform_action():
+            return
+
         self.dns_table[domain_name] = domain_ip_address
 
-    def reset_component_for_episode(self, episode: int):
-        """
-        Resets the Service component for a new episode.
-
-        This method ensures the Service is ready for a new episode, including resetting any
-        stateful properties or statistics, and clearing any message queues.
-        """
-        pass
-
     def receive(
         self,
         payload: Any,
@@ -88,10 +99,14 @@ class DNSServer(Service):
 
         :return: True if DNS request returns a valid IP, otherwise, False
         """
+        if not super().receive(payload=payload, session_id=session_id, **kwargs):
+            return False
+
         # The payload should be a DNS packet
         if not isinstance(payload, DNSPacket):
             _LOGGER.debug(f"{payload} is not a DNSPacket")
             return False
+
         # cast payload into a DNS packet
         payload: DNSPacket = payload
         if payload.dns_request is not None:

src/primaite/simulator/system/services/ftp/ftp_client.py

@@ -1,13 +1,15 @@
 from ipaddress import IPv4Address
 from typing import Optional
 
+from primaite import getLogger
 from primaite.simulator.file_system.file_system import File
 from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
 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.ftp.ftp_service import FTPServiceABC
-from primaite.simulator.system.services.service import ServiceOperatingState
+
+_LOGGER = getLogger(__name__)
 
 
 class FTPClient(FTPServiceABC):
@@ -28,6 +30,18 @@ class FTPClient(FTPServiceABC):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting FTPClient original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"connected"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting FTPClient state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
     def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
         """
         Process the command in the FTP Packet.
@@ -38,8 +52,7 @@ class FTPClient(FTPServiceABC):
         :type: session_id: Optional[str]
         """
         # if client service is down, return error
-        if self.operating_state != ServiceOperatingState.RUNNING:
-            self.sys_log.error("FTP Client is not running")
+        if not self._can_perform_action():
             payload.status_code = FTPStatusCode.ERROR
             return payload
 
@@ -66,16 +79,12 @@ class FTPClient(FTPServiceABC):
         :type: is_reattempt: Optional[bool]
         """
         # make sure the service is running before attempting
-        if self.operating_state != ServiceOperatingState.RUNNING:
-            self.sys_log.error(f"FTPClient not running for {self.sys_log.hostname}")
+        if not self._can_perform_action():
             return False
 
         # normally FTP will choose a random port for the transfer, but using the FTP command port will do for now
         # create FTP packet
-        payload: FTPPacket = FTPPacket(
-            ftp_command=FTPCommand.PORT,
-            ftp_command_args=Port.FTP,
-        )
+        payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.PORT, ftp_command_args=Port.FTP)
 
         if self.send(payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id):
             if payload.status_code == FTPStatusCode.OK:
@@ -270,8 +279,14 @@ class FTPClient(FTPServiceABC):
         This helps prevent an FTP request loop - FTP client and servers can exist on
         the same node.
         """
-        if payload.status_code is None:
+        if not self._can_perform_action():
             return False
 
+        if payload.status_code is None:
+            self.sys_log.error(f"FTP Server could not be found - Error Code: {FTPStatusCode.NOT_FOUND.value}")
+            return False
+
+        self.sys_log.info(f"{self.name}: Received FTP Response {payload.ftp_command.name} {payload.status_code.value}")
+
         self._process_ftp_command(payload=payload, session_id=session_id)
         return True

src/primaite/simulator/system/services/ftp/ftp_server.py

@@ -1,11 +1,13 @@
 from ipaddress import IPv4Address
 from typing import Any, Dict, Optional
 
+from primaite import getLogger
 from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
-from primaite.simulator.system.services.service import ServiceOperatingState
+
+_LOGGER = getLogger(__name__)
 
 
 class FTPServer(FTPServiceABC):
@@ -29,6 +31,19 @@ class FTPServer(FTPServiceABC):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting FTPServer original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"server_password"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting FTPServer state on node {self.software_manager.node.hostname}")
+        self.connections.clear()
+        super().reset_component_for_episode(episode)
+
     def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
         """
         Process the command in the FTP Packet.
@@ -42,8 +57,7 @@ class FTPServer(FTPServiceABC):
         payload.status_code = FTPStatusCode.ERROR
 
         # if server service is down, return error
-        if self.operating_state != ServiceOperatingState.RUNNING:
-            self.sys_log.error("FTP Server not running")
+        if not self._can_perform_action():
             return payload
 
         self.sys_log.info(f"{self.name}: Received FTP {payload.ftp_command.name} {payload.ftp_command_args}")
@@ -79,6 +93,9 @@ class FTPServer(FTPServiceABC):
             self.sys_log.error(f"{payload} is not an FTP packet")
             return False
 
+        if not super().receive(payload=payload, session_id=session_id, **kwargs):
+            return False
+
         """
         Ignore ftp payload if status code is defined.
 

src/primaite/simulator/system/services/red_services/data_manipulation_bot.py

@@ -1,27 +1,90 @@
+from enum import IntEnum
 from ipaddress import IPv4Address
 from typing import Optional
 
+from primaite import getLogger
+from primaite.game.science import simulate_trial
+from primaite.simulator.core import RequestManager, RequestType
+from primaite.simulator.system.applications.application import ApplicationOperatingState
 from primaite.simulator.system.applications.database_client import DatabaseClient
 
+_LOGGER = getLogger(__name__)
+
+
+class DataManipulationAttackStage(IntEnum):
+    """
+    Enumeration representing different stages of a data manipulation attack.
+
+    This enumeration defines the various stages a data manipulation attack can be in during its lifecycle in the
+    simulation. Each stage represents a specific phase in the attack process.
+    """
+
+    NOT_STARTED = 0
+    "Indicates that the attack has not started yet."
+    LOGON = 1
+    "The stage where logon procedures are simulated."
+    PORT_SCAN = 2
+    "Represents the stage of performing a horizontal port scan on the target."
+    ATTACKING = 3
+    "Stage of actively attacking the target."
+    SUCCEEDED = 4
+    "Indicates the attack has been successfully completed."
+    FAILED = 5
+    "Signifies that the attack has failed."
+
 
 class DataManipulationBot(DatabaseClient):
-    """
-    Red Agent Data Integration Service.
-
-    The Service represents a bot that causes files/folders in the File System to
-    become corrupted.
-    """
+    """A bot that simulates a script which performs a SQL injection attack."""
 
     server_ip_address: Optional[IPv4Address] = None
     payload: Optional[str] = None
     server_password: Optional[str] = None
+    port_scan_p_of_success: float = 0.1
+    data_manipulation_p_of_success: float = 0.1
+
+    attack_stage: DataManipulationAttackStage = DataManipulationAttackStage.NOT_STARTED
+    repeat: bool = False
+    "Whether to repeat attacking once finished."
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.name = "DataManipulationBot"
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting DataManipulationBot original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {
+            "server_ip_address",
+            "payload",
+            "server_password",
+            "port_scan_p_of_success",
+            "data_manipulation_p_of_success",
+            "attack_stage",
+            "repeat",
+        }
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting DataManipulationBot state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
+    def _init_request_manager(self) -> RequestManager:
+        rm = super()._init_request_manager()
+
+        rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
+
+        return rm
+
     def configure(
-        self, server_ip_address: IPv4Address, server_password: Optional[str] = None, payload: Optional[str] = None
+        self,
+        server_ip_address: IPv4Address,
+        server_password: Optional[str] = None,
+        payload: Optional[str] = None,
+        port_scan_p_of_success: float = 0.1,
+        data_manipulation_p_of_success: float = 0.1,
+        repeat: bool = False,
     ):
         """
         Configure the DataManipulatorBot to communicate with a DatabaseService.
@@ -29,23 +92,111 @@ class DataManipulationBot(DatabaseClient):
         :param server_ip_address: The IP address of the Node the DatabaseService is on.
         :param server_password: The password on the DatabaseService.
         :param payload: The data manipulation query payload.
+        :param port_scan_p_of_success: The probability of success for the port scan stage.
+        :param data_manipulation_p_of_success: The probability of success for the data manipulation stage.
+        :param repeat: Whether to repeat attacking once finished.
         """
         self.server_ip_address = server_ip_address
         self.payload = payload
         self.server_password = server_password
+        self.port_scan_p_of_success = port_scan_p_of_success
+        self.data_manipulation_p_of_success = data_manipulation_p_of_success
+        self.repeat = repeat
         self.sys_log.info(
-            f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}."
+            f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
+            f"{repeat=}."
         )
 
+    def _logon(self):
+        """
+        Simulate the logon process as the initial stage of the attack.
+
+        Advances the attack stage to `LOGON` if successful.
+        """
+        if self.attack_stage == DataManipulationAttackStage.NOT_STARTED:
+            # Bypass this stage as we're not dealing with logon for now
+            self.sys_log.info(f"{self.name}: ")
+            self.attack_stage = DataManipulationAttackStage.LOGON
+
+    def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
+        """
+        Perform a simulated port scan to check for open SQL ports.
+
+        Advances the attack stage to `PORT_SCAN` if successful.
+
+        :param p_of_success: Probability of successful port scan, by default 0.1.
+        """
+        if self.attack_stage == DataManipulationAttackStage.LOGON:
+            # perform a port scan to identify that the SQL port is open on the server
+            if simulate_trial(p_of_success):
+                self.sys_log.info(f"{self.name}: Performing port scan")
+                # perform the port scan
+                port_is_open = True  # Temporary; later we can implement NMAP port scan.
+                if port_is_open:
+                    self.sys_log.info(f"{self.name}: ")
+                    self.attack_stage = DataManipulationAttackStage.PORT_SCAN
+
+    def _perform_data_manipulation(self, p_of_success: Optional[float] = 0.1):
+        """
+        Execute the data manipulation attack on the target.
+
+        Advances the attack stage to `COMPLETE` if successful, or 'FAILED' if unsuccessful.
+
+        :param p_of_success: Probability of successfully performing data manipulation, by default 0.1.
+        """
+        if self.attack_stage == DataManipulationAttackStage.PORT_SCAN:
+            # perform the actual data manipulation attack
+            if simulate_trial(p_of_success):
+                self.sys_log.info(f"{self.name}: Performing data manipulation")
+                # perform the attack
+                if not self.connected:
+                    self.connect()
+                if self.connected:
+                    self.query(self.payload)
+                    self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
+                    attack_successful = True
+                    if attack_successful:
+                        self.sys_log.info(f"{self.name}: Data manipulation successful")
+                        self.attack_stage = DataManipulationAttackStage.SUCCEEDED
+                    else:
+                        self.sys_log.info(f"{self.name}: Data manipulation failed")
+                        self.attack_stage = DataManipulationAttackStage.FAILED
+
     def run(self):
-        """Run the DataManipulationBot."""
-        if self.server_ip_address and self.payload:
-            self.sys_log.info(f"{self.name}: Attempting to start the {self.name}")
-            super().run()
-            if not self.connected:
-                self.connect()
-            if self.connected:
-                self.query(self.payload)
-                self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
+        """
+        Run the Data Manipulation Bot.
+
+        Calls the parent classes execute method before starting the application loop.
+        """
+        super().run()
+        self._application_loop()
+
+    def _application_loop(self):
+        """
+        The main application loop of the bot, handling the attack process.
+
+        This is the core loop where the bot sequentially goes through the stages of the attack.
+        """
+        if self.operating_state != ApplicationOperatingState.RUNNING:
+            return
+        if self.server_ip_address and self.payload and self.operating_state:
+            self.sys_log.info(f"{self.name}: Running")
+            self._logon()
+            self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
+            self._perform_data_manipulation(p_of_success=self.data_manipulation_p_of_success)
+
+            if self.repeat and self.attack_stage in (
+                DataManipulationAttackStage.SUCCEEDED,
+                DataManipulationAttackStage.FAILED,
+            ):
+                self.attack_stage = DataManipulationAttackStage.NOT_STARTED
         else:
-            self.sys_log.error(f"Failed to start the {self.name} as it requires both a target_ip_address and payload.")
+            self.sys_log.error(f"{self.name}: Failed to start as it requires both a target_ip_address and payload.")
+
+    def apply_timestep(self, timestep: int) -> None:
+        """
+        Apply a timestep to the bot, triggering the application loop.
+
+        :param timestep: The timestep value to update the bot's state.
+        """
+        self._application_loop()

src/primaite/simulator/system/services/service.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from primaite import getLogger
 from primaite.simulator.core import RequestManager, RequestType
@@ -40,12 +40,52 @@ class Service(IOSoftware):
     restart_countdown: Optional[int] = None
     "If currently restarting, how many timesteps remain until the restart is finished."
 
+    def _can_perform_action(self) -> bool:
+        """
+        Checks if the service can perform actions.
+
+        This is done by checking if the service is operating properly or the node it is installed
+        in is operational.
+
+        Returns true if the software can perform actions.
+        """
+        if not super()._can_perform_action():
+            return False
+
+        if self.operating_state is not self.operating_state.RUNNING:
+            # service is not running
+            _LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
+            return False
+
+        return True
+
+    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
+        """
+        Receives a payload from the SessionManager.
+
+        The specifics of how the payload is processed and whether a response payload
+        is generated should be implemented in subclasses.
+
+
+        :param payload: The payload to receive.
+        :param session_id: The identifier of the session that the payload is associated with.
+        :param kwargs: Additional keyword arguments specific to the implementation.
+        :return: True if the payload was successfully received and processed, False otherwise.
+        """
+        return super().receive(payload=payload, session_id=session_id, **kwargs)
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
         self.health_state_visible = SoftwareHealthState.UNUSED
         self.health_state_actual = SoftwareHealthState.UNUSED
 
+    def set_original_state(self):
+        """Sets the original state."""
+        super().set_original_state()
+        vals_to_include = {"operating_state", "restart_duration", "restart_countdown"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
     def _init_request_manager(self) -> RequestManager:
         rm = super()._init_request_manager()
         rm.add_request("scan", RequestType(func=lambda request, context: self.scan()))
@@ -69,19 +109,10 @@ class Service(IOSoftware):
         """
         state = super().describe_state()
         state["operating_state"] = self.operating_state.value
-        state["health_state_actual"] = self.health_state_actual
-        state["health_state_visible"] = self.health_state_visible
+        state["health_state_actual"] = self.health_state_actual.value
+        state["health_state_visible"] = self.health_state_visible.value
         return state
 
-    def reset_component_for_episode(self, episode: int):
-        """
-        Resets the Service component for a new episode.
-
-        This method ensures the Service is ready for a new episode, including resetting any
-        stateful properties or statistics, and clearing any message queues.
-        """
-        pass
-
     def stop(self) -> None:
         """Stop the service."""
         if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
@@ -91,6 +122,10 @@ class Service(IOSoftware):
 
     def start(self, **kwargs) -> None:
         """Start the service."""
+        # cant start the service if the node it is on is off
+        if not super()._can_perform_action():
+            return
+
         if self.operating_state == ServiceOperatingState.STOPPED:
             self.sys_log.info(f"Starting service {self.name}")
             self.operating_state = ServiceOperatingState.RUNNING

src/primaite/simulator/system/services/web_server/web_server.py

@@ -2,6 +2,7 @@ from ipaddress import IPv4Address
 from typing import Any, Dict, Optional
 from urllib.parse import urlparse
 
+from primaite import getLogger
 from primaite.simulator.network.protocols.http import (
     HttpRequestMethod,
     HttpRequestPacket,
@@ -13,12 +14,26 @@ from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.system.applications.database_client import DatabaseClient
 from primaite.simulator.system.services.service import Service
 
+_LOGGER = getLogger(__name__)
+
 
 class WebServer(Service):
     """Class used to represent a Web Server Service in simulation."""
 
     last_response_status_code: Optional[HttpStatusCode] = None
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting WebServer original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"last_response_status_code"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting WebServer state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
     def describe_state(self) -> Dict:
         """
         Produce a dictionary describing the current state of this object.
@@ -30,7 +45,7 @@ class WebServer(Service):
         """
         state = super().describe_state()
         state["last_response_status_code"] = (
-            self.last_response_status_code.value if self.last_response_status_code else None
+            self.last_response_status_code.value if isinstance(self.last_response_status_code, HttpStatusCode) else None
         )
         return state
 
@@ -104,7 +119,7 @@ class WebServer(Service):
 
             if path.startswith("users"):
                 # get data from DatabaseServer
-                db_client: DatabaseClient = self.software_manager.software["DatabaseClient"]
+                db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
                 # get all users
                 if db_client.query("SELECT"):
                     # query succeeded
@@ -155,6 +170,9 @@ class WebServer(Service):
         :param: payload: The payload to send.
         :param: session_id: The id of the session. Optional.
         """
+        if not super().receive(payload=payload, session_id=session_id, **kwargs):
+            return False
+
         # check if the payload is an HTTPPacket
         if not isinstance(payload, HttpRequestPacket):
             self.sys_log.error("Payload is not an HTTPPacket")