diff --git a/src/primaite/simulator/file_system/file.py b/src/primaite/simulator/file_system/file.py index 8f0abb3c..f0984f89 100644 --- a/src/primaite/simulator/file_system/file.py +++ b/src/primaite/simulator/file_system/file.py @@ -77,12 +77,14 @@ class File(FileSystemItemABC): def set_original_state(self): """Sets the original state.""" + print(f"Setting File ({self.path}) original state on node {self.sys_log.hostname}") super().set_original_state() vals_to_include = {"folder_id", "folder_name", "file_type", "sim_size", "real", "sim_path", "sim_root"} 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.""" + print(f"Resetting File ({self.path}) state on node {self.sys_log.hostname}") super().reset_component_for_episode(episode) @property diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index dc6f01a3..a6876786 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -37,15 +37,20 @@ class FileSystem(SimComponent): def set_original_state(self): """Sets the original state.""" + print(f"Setting FileSystem original state on node {self.sys_log.hostname}") for folder in self.folders.values(): folder.set_original_state() - super().set_original_state() # Capture a list of all 'original' file uuids - self._original_state["original_folder_uuids"] = list(self.folders.keys()) + original_keys = list(self.folders.keys()) + vals_to_include = {"sim_root"} + self._original_state.update(self.model_dump(include=vals_to_include)) + self._original_state["original_folder_uuids"] = original_keys def reset_component_for_episode(self, episode: int): """Reset the original state of the SimComponent.""" + print(f"Resetting FileSystem state on node {self.sys_log.hostname}") # Move any 'original' folder that have been deleted back to folders + print(self._original_state) original_folder_uuids = self._original_state.pop("original_folder_uuids") for uuid in original_folder_uuids: if uuid in self.deleted_folders: diff --git a/src/primaite/simulator/file_system/folder.py b/src/primaite/simulator/file_system/folder.py index 8e577097..24dbdd79 100644 --- a/src/primaite/simulator/file_system/folder.py +++ b/src/primaite/simulator/file_system/folder.py @@ -53,6 +53,7 @@ class Folder(FileSystemItemABC): def set_original_state(self): """Sets the original state.""" + print(f"Setting Folder ({self.name}) original state on node {self.sys_log.hostname}") for file in self.files.values(): file.set_original_state() super().set_original_state() @@ -69,6 +70,7 @@ class Folder(FileSystemItemABC): def reset_component_for_episode(self, episode: int): """Reset the original state of the SimComponent.""" + print(f"Resetting Folder ({self.name}) state on node {self.sys_log.hostname}") # Move any 'original' file that have been deleted back to files original_file_uuids = self._original_state.pop("original_file_uuids") for uuid in original_file_uuids: diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index c6ee373e..b72fde54 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -993,6 +993,7 @@ class Node(SimComponent): def set_original_state(self): """Sets the original state.""" + print(f"Setting node original state for {self.hostname}") for software in self.software_manager.software.values(): software.set_original_state() @@ -1019,6 +1020,7 @@ class Node(SimComponent): def reset_component_for_episode(self, episode: int): """Reset the original state of the SimComponent.""" + print(f"Resetting node state for {self.hostname}") # Reset ARP Cache self.arp.clear() @@ -1031,6 +1033,10 @@ class Node(SimComponent): # Reset software for software in self.software_manager.software.values(): software.reset_component_for_episode(episode) + if isinstance(software, Service): + software.start() + elif isinstance(software, Application): + software.run() # Reset File System self.file_system.reset_component_for_episode(episode) @@ -1039,13 +1045,16 @@ class Node(SimComponent): for nic in self.nics.values(): nic.reset_component_for_episode(episode) - # if episode and self.sys_log: self.sys_log.current_episode = episode self.sys_log.setup_logger() super().reset_component_for_episode(episode) + self.power_on() + for nic in self.nics.values(): + nic.enable() + def _init_request_manager(self) -> RequestManager: # TODO: I see that this code is really confusing and hard to read right now... I think some of these things will # need a better name and better documentation. diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 34b92a07..0017215a 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -678,8 +678,9 @@ class Router(Node): """Sets the original state.""" self.acl.set_original_state() self.route_table.set_original_state() + super().set_original_state() vals_to_include = {"num_ports"} - self._original_state = self.model_dump(include=vals_to_include) + 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.""" diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index 37f85b28..92f7e76d 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -35,10 +35,16 @@ class DatabaseClient(Application): def set_original_state(self): """Sets the original state.""" + print(f"Setting DatabaseClient WebServer original state on node {self.software_manager.node.hostname}") super().set_original_state() vals_to_include = {"server_ip_address", "server_password", "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.""" + print(f"Resetting DataBaseClient state on node {self.software_manager.node.hostname}") + super().reset_component_for_episode(episode) + def describe_state(self) -> Dict: """ Describes the current state of the ACLRule. diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index bf304d7b..88560240 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -43,10 +43,16 @@ class WebBrowser(Application): def set_original_state(self): """Sets the original state.""" + print(f"Setting WebBrowser original state on node {self.software_manager.node.hostname}") super().set_original_state() vals_to_include = {"target_url", "domain_name_ip_address", "latest_response"} 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.""" + print(f"Resetting WebBrowser 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( diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 45e469fb..925d1df0 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -40,6 +40,7 @@ class DatabaseService(Service): def set_original_state(self): """Sets the original state.""" + print(f"Setting DatabaseService original state on node {self.software_manager.node.hostname}") super().set_original_state() vals_to_include = { "password", @@ -52,6 +53,7 @@ class DatabaseService(Service): def reset_component_for_episode(self, episode: int): """Reset the original state of the SimComponent.""" + print("Resetting DatabaseService original state on node {self.software_manager.node.hostname}") self.connections.clear() super().reset_component_for_episode(episode) diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 3d425bfa..147387ae 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/src/primaite/simulator/system/services/dns/dns_client.py @@ -31,12 +31,14 @@ class DNSClient(Service): def set_original_state(self): """Sets the original state.""" + print(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.""" + print(f"Resetting DNSClient state on node {self.software_manager.node.hostname}") self.dns_cache.clear() super().reset_component_for_episode(episode) @@ -53,15 +55,6 @@ 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) -> bool: """ Adds a domain name to the DNS Client cache. diff --git a/src/primaite/simulator/system/services/dns/dns_server.py b/src/primaite/simulator/system/services/dns/dns_server.py index 30278ab1..7842a07e 100644 --- a/src/primaite/simulator/system/services/dns/dns_server.py +++ b/src/primaite/simulator/system/services/dns/dns_server.py @@ -30,19 +30,18 @@ class DNSServer(Service): def set_original_state(self): """Sets the original state.""" + print(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.""" - print("dns reset") - print("DNSServer original state", self._original_state) + print(f"Resetting DNSServer state on node {self.software_manager.node.hostname}") 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) - self.show() def describe_state(self) -> Dict: """ diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 649b9b50..011b597f 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -28,6 +28,18 @@ class FTPClient(FTPServiceABC): super().__init__(**kwargs) self.start() + def set_original_state(self): + """Sets the original state.""" + print(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.""" + print(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. diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index cd128339..811a8939 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -29,6 +29,19 @@ class FTPServer(FTPServiceABC): super().__init__(**kwargs) self.start() + def set_original_state(self): + """Sets the original state.""" + print(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.""" + print(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. diff --git a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py index b0b34396..75cdee85 100644 --- a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py +++ b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py @@ -47,6 +47,26 @@ class DataManipulationBot(DatabaseClient): super().__init__(**kwargs) self.name = "DataManipulationBot" + def set_original_state(self): + """Sets the original state.""" + print(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.""" + print(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() diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py index becbf9f9..f34bba37 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/src/primaite/simulator/system/services/web_server/web_server.py @@ -17,22 +17,20 @@ from primaite.simulator.system.services.service import Service class WebServer(Service): """Class used to represent a Web Server Service in simulation.""" - _last_response_status_code: Optional[HttpStatusCode] = None + last_response_status_code: Optional[HttpStatusCode] = None + + def set_original_state(self): + """Sets the original state.""" + print(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.""" - self._last_response_status_code = None + print(f"Resetting WebServer state on node {self.software_manager.node.hostname}") super().reset_component_for_episode(episode) - @property - def last_response_status_code(self) -> HttpStatusCode: - """The latest http response code.""" - return self._last_response_status_code - - @last_response_status_code.setter - def last_response_status_code(self, val: Any): - self._last_response_status_code = val - def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object.