diff --git a/sandbox.py b/sandbox.py new file mode 100644 index 00000000..b08f15b1 --- /dev/null +++ b/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() diff --git a/src/primaite/simulator/file_system/file.py b/src/primaite/simulator/file_system/file.py index f0984f89..608a1d78 100644 --- a/src/primaite/simulator/file_system/file.py +++ b/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 diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index a6876786..31a3c5a0 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/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: diff --git a/src/primaite/simulator/file_system/folder.py b/src/primaite/simulator/file_system/folder.py index 24dbdd79..c45dd8c5 100644 --- a/src/primaite/simulator/file_system/folder.py +++ b/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: diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 1ee384f3..97b62f95 100644 --- a/src/primaite/simulator/network/container.py +++ b/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() diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 9fb007ff..04c76c6b 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/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. diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index b1743fad..7b63d26e 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/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 diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index 88560240..8f12df4e 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/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: diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 925d1df0..f9621ba5 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/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", diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 147387ae..2d3879ff 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/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) diff --git a/src/primaite/simulator/system/services/dns/dns_server.py b/src/primaite/simulator/system/services/dns/dns_server.py index 7842a07e..8decf7e9 100644 --- a/src/primaite/simulator/system/services/dns/dns_server.py +++ b/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 diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 011b597f..23d52342 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/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: diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index 811a8939..44d0455f 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/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) 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 75cdee85..44a56cf1 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 @@ -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: 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 f34bba37..bff29a47 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/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: