Fix software reset issues

sandbox.py

@@ -0,0 +1,72 @@
+from primaite.config.load import example_config_path, load
+from primaite.session.session import PrimaiteSession
+from primaite.simulator.system.applications.database_client import DatabaseClient
+from primaite.simulator.system.applications.web_browser import WebBrowser
+from primaite.simulator.system.services.dns.dns_client import DNSClient
+
+cfg = load(example_config_path())
+session = PrimaiteSession.from_config(cfg)
+network = session.game.simulation.network
+
+dc = network.get_node_by_hostname("domain_controller")
+router = network.get_node_by_hostname("router_1")
+client_1 = network.get_node_by_hostname("client_1")
+client_2 = network.get_node_by_hostname("client_2")
+switch_1 = network.get_node_by_hostname("switch_1")
+switch_2 = network.get_node_by_hostname("switch_2")
+web_server = network.get_node_by_hostname("web_server")
+
+dns_server = dc.software_manager.software["DNSServer"]
+dns_client: DNSClient = client_2.software_manager.software["DNSClient"]
+web_db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+web_browser: WebBrowser = client_2.software_manager.software["WebBrowser"]
+
+# print("before calling get webpage")
+# router.acl.show()
+# dns_server.show()
+# client_2.arp.show()
+# router.arp.show()
+# print()
+
+# print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+# print("after calling get webpage")
+# router.acl.show()
+# dns_server.show()
+# client_2.arp.show()
+# router.arp.show()
+# print()
+# print("reset")
+# print()
+# print("im gonna reset")
+# print()
+
+# web_db_client.connect()
+# web_db_client.run()
+# web_browser.run()
+# print("client_2", client_2.operating_state)
+# print("web_browser", web_browser.operating_state)
+# print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+session.game.reset()
+print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+session.game.reset()
+print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+session.game.reset()
+print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+session.game.reset()
+print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+# print()
+#
+# print("before calling get webpage")
+# router.acl.show()
+# dns_server.show()
+# client_2.arp.show()
+# router.arp.show()
+# print()
+#
+# print("can get webpage", client_2.software_manager.software["WebBrowser"].get_webpage())
+# print("after calling get webpage")
+# router.acl.show()
+# dns_server.show()
+# client_2.arp.show()
+# router.arp.show()
+# print()

src/primaite/simulator/file_system/file.py

@@ -77,14 +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}")
+        _LOGGER.debug(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}")
+        _LOGGER.debug(f"Resetting File ({self.path}) state on node {self.sys_log.hostname}")
         super().reset_component_for_episode(episode)
 
     @property

src/primaite/simulator/file_system/file_system.py

@@ -37,7 +37,7 @@ class FileSystem(SimComponent):
 
     def set_original_state(self):
         """Sets the original state."""
-        print(f"Setting FileSystem original state on node {self.sys_log.hostname}")
+        _LOGGER.debug(f"Setting FileSystem original state on node {self.sys_log.hostname}")
         for folder in self.folders.values():
             folder.set_original_state()
         # Capture a list of all 'original' file uuids
@@ -48,9 +48,8 @@ class FileSystem(SimComponent):
 
     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}")
+        _LOGGER.debug(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:

src/primaite/simulator/file_system/folder.py

@@ -53,7 +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}")
+        _LOGGER.debug(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()
@@ -70,7 +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}")
+        _LOGGER.debug(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:

src/primaite/simulator/network/container.py

@@ -58,21 +58,19 @@ class Network(SimComponent):
             node.reset_component_for_episode(episode)
         for link in self.links.values():
             link.reset_component_for_episode(episode)
-        
+
         for node in self.nodes.values():
             node.power_on()
 
+            for nic in node.nics.values():
+                nic.enable()
             # Reset software
             for software in node.software_manager.software.values():
-                software.reset_component_for_episode(episode)
                 if isinstance(software, Service):
                     software.start()
                 elif isinstance(software, Application):
                     software.run()
 
-            for nic in node.nics.values():
-                nic.enable()
-
     def _init_request_manager(self) -> RequestManager:
         rm = super()._init_request_manager()
         self._node_request_manager = RequestManager()

src/primaite/simulator/network/hardware/base.py

@@ -993,7 +993,6 @@ 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()
 
@@ -1020,7 +1019,6 @@ 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}")
         super().reset_component_for_episode(episode)
 
         # Reset ARP Cache
@@ -1039,11 +1037,13 @@ class Node(SimComponent):
         for nic in self.nics.values():
             nic.reset_component_for_episode(episode)
 
+        for software in self.software_manager.software.values():
+            software.reset_component_for_episode(episode)
+
         if episode and self.sys_log:
             self.sys_log.current_episode = episode
             self.sys_log.setup_logger()
 
-
     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.

src/primaite/simulator/system/applications/database_client.py

@@ -35,14 +35,14 @@ 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}")
+        _LOGGER.debug(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", "_query_success_tracker"}
         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}")
+        _LOGGER.debug(f"Resetting DataBaseClient state on node {self.software_manager.node.hostname}")
         super().reset_component_for_episode(episode)
         self._query_success_tracker.clear()
 
@@ -195,4 +195,6 @@ class DatabaseClient(Application):
                 self._query_success_tracker[query_id] = status_code == 200
                 if self._query_success_tracker[query_id]:
                     _LOGGER.debug(f"Received payload {payload}")
+                else:
+                    self.connected = False
         return True

src/primaite/simulator/system/applications/web_browser.py

@@ -2,6 +2,7 @@ from ipaddress import IPv4Address
 from typing import Dict, Optional
 from urllib.parse import urlparse
 
+from primaite import getLogger
 from primaite.simulator.core import RequestManager, RequestType
 from primaite.simulator.network.protocols.http import (
     HttpRequestMethod,
@@ -14,6 +15,8 @@ from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.services.dns.dns_client import DNSClient
 
+_LOGGER = getLogger(__name__)
+
 
 class WebBrowser(Application):
     """
@@ -43,14 +46,14 @@ 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}")
+        _LOGGER.debug(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}")
+        _LOGGER.debug(f"Resetting WebBrowser state on node {self.software_manager.node.hostname}")
         super().reset_component_for_episode(episode)
 
     def _init_request_manager(self) -> RequestManager:

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):
     """
@@ -40,7 +43,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}")
+        _LOGGER.debug(f"Setting DatabaseService original state on node {self.software_manager.node.hostname}")
         super().set_original_state()
         vals_to_include = {
             "password",

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

@@ -31,14 +31,13 @@ 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}")
+        _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."""
-        print(f"Resetting DNSClient state on node {self.software_manager.node.hostname}")
         self.dns_cache.clear()
         super().reset_component_for_episode(episode)
 

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

@@ -30,14 +30,13 @@ 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}")
+        _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."""
-        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

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

@@ -1,6 +1,7 @@
 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
@@ -9,6 +10,8 @@ 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):
     """
@@ -30,14 +33,14 @@ class FTPClient(FTPServiceABC):
 
     def set_original_state(self):
         """Sets the original state."""
-        print(f"Setting FTPClient original state on node {self.software_manager.node.hostname}")
+        _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."""
-        print(f"Resetting FTPClient state on node {self.software_manager.node.hostname}")
+        _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:

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

@@ -1,12 +1,15 @@
 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):
     """
@@ -31,14 +34,14 @@ class FTPServer(FTPServiceABC):
 
     def set_original_state(self):
         """Sets the original state."""
-        print(f"Setting FTPServer original state on node {self.software_manager.node.hostname}")
+        _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."""
-        print(f"Resetting FTPServer state on node {self.software_manager.node.hostname}")
+        _LOGGER.debug(f"Resetting FTPServer state on node {self.software_manager.node.hostname}")
         self.connections.clear()
         super().reset_component_for_episode(episode)
 

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

@@ -2,11 +2,14 @@ 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):
     """
@@ -49,7 +52,7 @@ class DataManipulationBot(DatabaseClient):
 
     def set_original_state(self):
         """Sets the original state."""
-        print(f"Setting DataManipulationBot original state on node {self.software_manager.node.hostname}")
+        _LOGGER.debug(f"Setting DataManipulationBot original state on node {self.software_manager.node.hostname}")
         super().set_original_state()
         vals_to_include = {
             "server_ip_address",
@@ -64,7 +67,7 @@ class DataManipulationBot(DatabaseClient):
 
     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}")
+        _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:

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,6 +14,8 @@ 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."""
@@ -21,14 +24,14 @@ class WebServer(Service):
 
     def set_original_state(self):
         """Sets the original state."""
-        print(f"Setting WebServer original state on node {self.software_manager.node.hostname}")
+        _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."""
-        print(f"Resetting WebServer state on node {self.software_manager.node.hostname}")
+        _LOGGER.debug(f"Resetting WebServer state on node {self.software_manager.node.hostname}")
         super().reset_component_for_episode(episode)
 
     def describe_state(self) -> Dict: