diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 5da4eca8..5653234d 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -3,6 +3,8 @@ from __future__ import annotations import math import os.path import shutil +from abc import abstractmethod +from enum import Enum from pathlib import Path from typing import Dict, Optional @@ -42,6 +44,19 @@ def convert_size(size_bytes: int) -> str: return f"{s} {size_name[i]}" +class FileSystemItemStatus(Enum): + """Status of the FileSystemItem.""" + + GOOD = 0 + """File/Folder is OK.""" + + QUARANTINED = 1 + """File/Folder is quarantined.""" + + CORRUPTED = 2 + """File/Folder is corrupted.""" + + class FileSystemItemABC(SimComponent): """ Abstract base class for file system items used in the file system simulation. @@ -52,6 +67,12 @@ class FileSystemItemABC(SimComponent): name: str "The name of the FileSystemItemABC." + status: FileSystemItemStatus = FileSystemItemStatus.GOOD + "Actual status of the current FileSystemItem" + + visible_status: FileSystemItemStatus = FileSystemItemStatus.GOOD + "Visible status of the current FileSystemItem" + def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -78,6 +99,32 @@ class FileSystemItemABC(SimComponent): """ return convert_size(self.size) + def _init_request_manager(self) -> RequestManager: + am = super()._init_request_manager() + am.add_request("scan", RequestType(func=lambda request, context: self.scan())) # TODO implement request + am.add_request("checkhash", RequestType(func=lambda request, context: ...)) # TODO implement request + am.add_request("delete", RequestType(func=lambda request, context: ...)) # TODO implement request + am.add_request("restore", RequestType(func=lambda request, context: ...)) # TODO implement request + am.add_request("repair", RequestType(func=lambda request, context: self.repair())) + am.add_request("corrupt", RequestType(func=lambda request, context: self.corrupt())) + return am + + def scan(self) -> None: + """Update the FileSystemItem states.""" + super().scan() + + self.visible_status = self.status + + @abstractmethod + def repair(self) -> None: + """Repair the FileSystemItem.""" + pass + + @abstractmethod + def corrupt(self) -> None: + """Corrupt the FileSystemItem.""" + pass + class FileSystem(SimComponent): """Class that contains all the simulation File System.""" @@ -329,20 +376,6 @@ class Folder(FileSystemItemABC): "Files stored in the folder." _files_by_name: Dict[str, File] = {} "Files by their name as .." - is_quarantined: bool = False - "Flag that marks the folder as quarantined if true." - - def _init_request_manager(self) -> RequestManager: - am = super()._init_request_manager() - - am.add_request("scan", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("checkhash", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("repair", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("restore", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("delete", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("corrupt", RequestType(func=lambda request, context: ...)) # TODO implement request - - return am def describe_state(self) -> Dict: """ @@ -440,19 +473,49 @@ class Folder(FileSystemItemABC): def quarantine(self): """Quarantines the File System Folder.""" - if not self.is_quarantined: - self.is_quarantined = True + if self.status != FileSystemItemStatus.QUARANTINED: + self.status = FileSystemItemStatus.QUARANTINED self.fs.sys_log.info(f"Quarantined folder ./{self.name}") def unquarantine(self): """Unquarantine of the File System Folder.""" - if self.is_quarantined: - self.is_quarantined = False + if self.status == FileSystemItemStatus.QUARANTINED: + self.status = FileSystemItemStatus.GOOD self.fs.sys_log.info(f"Quarantined folder ./{self.name}") def quarantine_status(self) -> bool: """Returns true if the folder is being quarantined.""" - return self.is_quarantined + return self.status == FileSystemItemStatus.QUARANTINED + + def repair(self) -> None: + """Repair a corrupted Folder by setting the folder and containing files status to FileSystemItemStatus.GOOD.""" + super().repair() + + # iterate through the files in the folder + for file_id in self.files: + file = self.get_file_by_id(file_uuid=file_id) + file.repair() + + # set file status to good if corrupt + if self.status == FileSystemItemStatus.CORRUPTED: + self.status = FileSystemItemStatus.GOOD + + self.fs.sys_log.info(f"Repaired folder {self.name}") + + def corrupt(self) -> None: + """Corrupt a File by setting the folder and containing files status to FileSystemItemStatus.CORRUPTED.""" + super().corrupt() + + # iterate through the files in the folder + for file_id in self.files: + file = self.get_file_by_id(file_uuid=file_id) + file.corrupt() + + # set file status to good if corrupt + if self.status == FileSystemItemStatus.GOOD: + self.status = FileSystemItemStatus.CORRUPTED + + self.fs.sys_log.info(f"Corrupted folder {self.name}") class File(FileSystemItemABC): @@ -509,18 +572,6 @@ class File(FileSystemItemABC): with open(self.sim_path, mode="a"): pass - def _init_request_manager(self) -> RequestManager: - am = super()._init_request_manager() - - am.add_request("scan", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("checkhash", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("delete", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("repair", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("restore", RequestType(func=lambda request, context: ...)) # TODO implement request - am.add_request("corrupt", RequestType(func=lambda request, context: ...)) # TODO implement request - - return am - def make_copy(self, dst_folder: Folder) -> File: """ Create a copy of the current File object in the given destination folder. @@ -556,3 +607,25 @@ class File(FileSystemItemABC): state["size"] = self.size state["file_type"] = self.file_type.name return state + + def repair(self) -> None: + """Repair a corrupted File by setting the status to FileSystemItemStatus.GOOD.""" + super().repair() + + # set file status to good if corrupt + if self.status == FileSystemItemStatus.CORRUPTED: + self.status = FileSystemItemStatus.GOOD + + path = self.folder.name + "/" + self.name + self.folder.fs.sys_log.info(f"Repaired file {self.sim_path if self.sim_path else path}") + + def corrupt(self) -> None: + """Corrupt a File by setting the status to FileSystemItemStatus.CORRUPTED.""" + super().corrupt() + + # set file status to good if corrupt + if self.status == FileSystemItemStatus.GOOD: + self.status = FileSystemItemStatus.CORRUPTED + + path = self.folder.name + "/" + self.name + self.folder.fs.sys_log.info(f"Corrupted file {self.sim_path if self.sim_path else path}") diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 597c8cbd..ca09ae60 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -40,7 +40,7 @@ class Service(IOSoftware): restart_duration: int = 5 "How many timesteps does it take to restart this service." - _restart_countdown: Optional[int] = None + restart_countdown: Optional[int] = None "If currently restarting, how many timesteps remain until the restart is finished." def _init_request_manager(self) -> RequestManager: @@ -65,7 +65,9 @@ class Service(IOSoftware): :rtype: Dict """ state = super().describe_state() - state.update({"operating_state": self.operating_state.name}) + state.update( + {"operating_state": self.operating_state.name, "visible_operating_state": self.visible_operating_state.name} + ) return state def reset_component_for_episode(self, episode: int): @@ -114,7 +116,7 @@ class Service(IOSoftware): if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: self.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.RESTARTING - self.restart_countdown = self.restarting_duration + self.restart_countdown = self.restart_duration def disable(self) -> None: """Disable the service.""" diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 33a03c1c..1fafd137 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -173,6 +173,8 @@ class Software(SimComponent): def scan(self) -> None: """Update the observed health status to match the actual health status.""" + super().scan() + self.health_state_visible = self.health_state_actual diff --git a/tests/conftest.py b/tests/conftest.py index 35548f2a..06e55400 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,7 @@ from primaite.environment.primaite_env import Primaite from primaite.primaite_session import PrimaiteSession from primaite.simulator.network.container import Network from primaite.simulator.network.networks import arcd_uc2_network +from primaite.simulator.system.core.sys_log import SysLog from tests.mock_and_patch.get_session_path_mock import get_temp_session_path ACTION_SPACE_NODE_VALUES = 1 diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py index d1d78003..539f2874 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py @@ -1,6 +1,6 @@ import pytest -from primaite.simulator.file_system.file_system import FileSystem +from primaite.simulator.file_system.file_system import File, FileSystem, FileSystemItemStatus, Folder from primaite.simulator.file_system.file_type import FileType @@ -135,6 +135,40 @@ def test_folder_quarantine_state(file_system): assert folder.quarantine_status() is False +def test_file_corrupt_repair(file_system): + """Test the ability to corrupt and repair files.""" + folder: Folder = file_system.create_folder(folder_name="test_folder") + file: File = file_system.create_file(file_name="test_file.txt", folder_name="test_folder") + + file.corrupt() + + assert folder.status == FileSystemItemStatus.GOOD + assert file.status == FileSystemItemStatus.CORRUPTED + + file.repair() + + assert folder.status == FileSystemItemStatus.GOOD + assert file.status == FileSystemItemStatus.GOOD + + +def test_folder_corrupt_repair(file_system): + """Test the ability to corrupt and repair folders.""" + folder: Folder = file_system.create_folder(folder_name="test_folder") + file_system.create_file(file_name="test_file.txt", folder_name="test_folder") + + folder.corrupt() + + file = folder.get_file(file_name="test_file.txt") + assert folder.status == FileSystemItemStatus.CORRUPTED + assert file.status == FileSystemItemStatus.CORRUPTED + + folder.repair() + + file = folder.get_file(file_name="test_file.txt") + assert folder.status == FileSystemItemStatus.GOOD + assert file.status == FileSystemItemStatus.GOOD + + @pytest.mark.skip(reason="Skipping until we tackle serialisation") def test_serialisation(file_system): """Test to check that the object serialisation works correctly.""" diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py new file mode 100644 index 00000000..20a3cad5 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py @@ -0,0 +1,80 @@ +from typing import Any + +import pytest + +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.core.sys_log import SysLog +from primaite.simulator.system.services.service import Service, ServiceOperatingState + + +class TestService(Service): + """Test Service class""" + + def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + pass + + +@pytest.fixture(scope="function") +def service(file_system) -> TestService: + return TestService( + name="TestService", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="test_service") + ) + + +def test_scan(service): + assert service.operating_state == ServiceOperatingState.STOPPED + assert service.visible_operating_state == ServiceOperatingState.STOPPED + + service.start() + assert service.operating_state == ServiceOperatingState.RUNNING + assert service.visible_operating_state == ServiceOperatingState.STOPPED + + service.scan() + assert service.operating_state == ServiceOperatingState.RUNNING + assert service.visible_operating_state == ServiceOperatingState.RUNNING + + +def test_start_service(service): + assert service.operating_state == ServiceOperatingState.STOPPED + service.start() + + assert service.operating_state == ServiceOperatingState.RUNNING + + +def test_stop_service(service): + service.start() + assert service.operating_state == ServiceOperatingState.RUNNING + + service.stop() + assert service.operating_state == ServiceOperatingState.STOPPED + + +def test_pause_and_resume_service(service): + assert service.operating_state == ServiceOperatingState.STOPPED + service.resume() + assert service.operating_state == ServiceOperatingState.STOPPED + + service.start() + service.pause() + assert service.operating_state == ServiceOperatingState.PAUSED + + service.resume() + assert service.operating_state == ServiceOperatingState.RUNNING + + +def test_restart(service): + assert service.operating_state == ServiceOperatingState.STOPPED + service.restart() + assert service.operating_state == ServiceOperatingState.STOPPED + + service.start() + service.restart() + assert service.operating_state == ServiceOperatingState.RESTARTING + + +def test_enable_disable(service): + service.disable() + assert service.operating_state == ServiceOperatingState.DISABLED + + service.enable() + assert service.operating_state == ServiceOperatingState.STOPPED