From a4d372d3ebe556fca63ebf6beb7f4d6f55c7fe60 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 9 Jan 2024 16:29:40 +0000 Subject: [PATCH] #2151: utilise set_health_state method instead of directly changing software states --- .../system/applications/application.py | 5 +--- .../services/database/database_service.py | 2 +- .../simulator/system/services/service.py | 2 +- src/primaite/simulator/system/software.py | 4 +-- .../test_dos_bot_and_server.py | 4 +-- .../_network/_hardware/test_node_actions.py | 6 ++-- .../_system/_services/test_services.py | 15 ++++------ .../_simulator/_system/test_software.py | 29 +++++++++++++++++++ 8 files changed, 45 insertions(+), 22 deletions(-) create mode 100644 tests/unit_tests/_primaite/_simulator/_system/test_software.py diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 898e5917..09828b89 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Dict, Set from primaite import getLogger -from primaite.simulator.system.software import IOSoftware, SoftwareHealthState +from primaite.simulator.system.software import IOSoftware _LOGGER = getLogger(__name__) @@ -38,9 +38,6 @@ class Application(IOSoftware): 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() diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 6f333091..1df1db9e 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -196,7 +196,7 @@ class DatabaseService(Service): return {"status_code": 404, "data": False} elif query == "DELETE": if self.health_state_actual == SoftwareHealthState.GOOD: - self.health_state_actual = SoftwareHealthState.COMPROMISED + self.set_health_state(SoftwareHealthState.COMPROMISED) return { "status_code": 200, "type": "sql", diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 1de52e92..43c85471 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -127,7 +127,7 @@ class Service(IOSoftware): self.operating_state = ServiceOperatingState.RUNNING # set software health state to GOOD if initially set to UNUSED if self.health_state_actual == SoftwareHealthState.UNUSED: - self.health_state_actual = SoftwareHealthState.GOOD + self.set_health_state(SoftwareHealthState.GOOD) def pause(self) -> None: """Pause the service.""" diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index f41a5a86..4072fab1 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -303,13 +303,13 @@ class IOSoftware(Software): """ # if over or at capacity, set to overwhelmed if len(self._connections) >= self.max_sessions: - self.health_state_actual = SoftwareHealthState.OVERWHELMED + self.set_health_state(SoftwareHealthState.OVERWHELMED) self.sys_log.error(f"{self.name}: Connect request for {connection_id=} declined. Service is at capacity.") return False else: # if service was previously overwhelmed, set to good because there is enough space for connections if self.health_state_actual == SoftwareHealthState.OVERWHELMED: - self.health_state_actual = SoftwareHealthState.GOOD + self.set_health_state(SoftwareHealthState.GOOD) # check that connection already doesn't exist if not self._connections.get(connection_id): diff --git a/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py b/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py index 85028d75..fb768127 100644 --- a/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py +++ b/tests/integration_tests/system/red_applications/test_dos_bot_and_server.py @@ -90,7 +90,7 @@ def test_repeating_dos_attack(dos_bot_and_db_server): assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED db_server_service.clear_connections() - db_server_service.health_state_actual = SoftwareHealthState.GOOD + db_server_service.set_health_state(SoftwareHealthState.GOOD) assert len(db_server_service.connections) == 0 computer.apply_timestep(timestep=1) @@ -121,7 +121,7 @@ def test_non_repeating_dos_attack(dos_bot_and_db_server): assert db_server_service.health_state_actual is SoftwareHealthState.OVERWHELMED db_server_service.clear_connections() - db_server_service.health_state_actual = SoftwareHealthState.GOOD + db_server_service.set_health_state(SoftwareHealthState.GOOD) assert len(db_server_service.connections) == 0 computer.apply_timestep(timestep=1) diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py index 5fe5df16..b6f7a86d 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py @@ -53,12 +53,12 @@ def test_node_os_scan(node, service, application): # TODO implement processes # add services to node - service.health_state_actual = SoftwareHealthState.COMPROMISED + service.set_health_state(SoftwareHealthState.COMPROMISED) node.install_service(service=service) assert service.health_state_visible == SoftwareHealthState.UNUSED # add application to node - application.health_state_actual = SoftwareHealthState.COMPROMISED + application.set_health_state(SoftwareHealthState.COMPROMISED) node.install_application(application=application) assert application.health_state_visible == SoftwareHealthState.UNUSED @@ -101,7 +101,7 @@ def test_node_red_scan(node, service, application): assert service.revealed_to_red is False # add application to node - application.health_state_actual = SoftwareHealthState.COMPROMISED + application.set_health_state(SoftwareHealthState.COMPROMISED) node.install_application(application=application) assert application.revealed_to_red is False diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py index 2c0671d5..ac36c660 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py @@ -57,12 +57,15 @@ def test_restart(service): assert service.operating_state == ServiceOperatingState.STOPPED assert service.health_state_actual == SoftwareHealthState.UNUSED service.restart() + # Service is STOPPED. Restart will only work if the service was PAUSED or RUNNING assert service.operating_state == ServiceOperatingState.STOPPED assert service.health_state_actual == SoftwareHealthState.UNUSED service.start() + assert service.operating_state == ServiceOperatingState.RUNNING assert service.health_state_actual == SoftwareHealthState.GOOD service.restart() + # Service is RUNNING. Restart should work assert service.operating_state == ServiceOperatingState.RESTARTING assert service.health_state_actual == SoftwareHealthState.GOOD @@ -77,17 +80,11 @@ def test_restart(service): def test_restart_compromised(service): - assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED - service.restart() - assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED - service.start() assert service.health_state_actual == SoftwareHealthState.GOOD # compromise the service - service.health_state_actual = SoftwareHealthState.COMPROMISED + service.set_health_state(SoftwareHealthState.COMPROMISED) service.restart() assert service.operating_state == ServiceOperatingState.RESTARTING @@ -118,7 +115,7 @@ def test_compromised_service_remains_compromised(service): service.start() assert service.health_state_actual == SoftwareHealthState.GOOD - service.health_state_actual = SoftwareHealthState.COMPROMISED + service.set_health_state(SoftwareHealthState.COMPROMISED) service.stop() assert service.health_state_actual == SoftwareHealthState.COMPROMISED @@ -143,7 +140,7 @@ def test_service_patching(service): service.start() assert service.health_state_actual == SoftwareHealthState.GOOD - service.health_state_actual = SoftwareHealthState.COMPROMISED + service.set_health_state(SoftwareHealthState.COMPROMISED) service.patch() assert service.health_state_actual == SoftwareHealthState.PATCHING diff --git a/tests/unit_tests/_primaite/_simulator/_system/test_software.py b/tests/unit_tests/_primaite/_simulator/_system/test_software.py new file mode 100644 index 00000000..e77cd895 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_system/test_software.py @@ -0,0 +1,29 @@ +from typing import Dict + +import pytest + +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.core.sys_log import SysLog +from primaite.simulator.system.software import Software, SoftwareHealthState + + +class TestSoftware(Software): + def describe_state(self) -> Dict: + pass + + +@pytest.fixture(scope="function") +def software(file_system): + return TestSoftware( + name="TestSoftware", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="test_service") + ) + + +def test_software_creation(software): + assert software is not None + + +def test_software_set_health_state(software): + assert software.health_state_actual == SoftwareHealthState.UNUSED + software.set_health_state(SoftwareHealthState.GOOD) + assert software.health_state_actual == SoftwareHealthState.GOOD