oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#1859 - Added the call to file system reset

Browse Source
This commit is contained in:
Chris McCarthy
2023-11-28 11:58:09 +00:00
parent da955038f5
commit 9a4855e7bd
14 changed files with 95 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/file_system/file.py
View File

@@ -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

9
src/primaite/simulator/file_system/file_system.py
View File

@@ -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:

2
src/primaite/simulator/file_system/folder.py
View File

@@ -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:

11
src/primaite/simulator/network/hardware/base.py
View File

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

3
src/primaite/simulator/network/hardware/nodes/router.py
View File

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

6
src/primaite/simulator/system/applications/database_client.py
View File

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

6
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -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(

2
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -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)

11
src/primaite/simulator/system/services/dns/dns_client.py
View File

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

5
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -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:
"""

12
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

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

13
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

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

20
src/primaite/simulator/system/services/red_services/data_manipulation_bot.py
View File

@@ -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()

20
src/primaite/simulator/system/services/web_server/web_server.py
View File

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