#1947: added ability for files and folders to be scanned, corrupted and repaired

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 <file name>.<file type>."
-    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:
+        if self.status != FileSystemItemStatus.QUARANTINED:
-            self.is_quarantined = True
+            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:
+        if self.status == FileSystemItemStatus.QUARANTINED:
-            self.is_quarantined = False
+            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}")

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."""

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
 
 

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

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."""

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