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

Merge branch 'feature/2087_Full-reset-functionality' into feature/2085-dump_describe_state

Browse Source
This commit is contained in:
Nick Todd
2023-11-30 16:16:44 +00:00
parent 5cd69f343f a16116a688
commit 10d4b61c73
15 changed files with 134 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
sandbox.py Normal file
View File

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

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

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

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

@@ -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,10 +48,9 @@ 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")
original_folder_uuids = self._original_state["original_folder_uuids"]
for uuid in original_folder_uuids:
if uuid in self.deleted_folders:
self.folders[uuid] = self.deleted_folders.pop(uuid)

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

@@ -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,9 +70,9 @@ 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")
original_file_uuids = self._original_state["original_file_uuids"]
for uuid in original_file_uuids:
if uuid in self.deleted_files:
self.files[uuid] = self.deleted_files.pop(uuid)

14
src/primaite/simulator/network/container.py
View File

@@ -12,6 +12,8 @@ from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import Router
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service
_LOGGER = getLogger(__name__)
@@ -57,6 +59,18 @@ class Network(SimComponent):
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():
if isinstance(software, Service):
software.start()
elif isinstance(software, Application):
software.run()
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
self._node_request_manager = RequestManager()

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

@@ -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
@@ -1035,27 +1033,17 @@ class Node(SimComponent):
# Reset File System
self.file_system.reset_component_for_episode(episode)
self.power_on()
# Reset all Nics
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()
# 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()
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.

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

@@ -35,15 +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}")
_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"}
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()
def describe_state(self) -> Dict:
"""
@@ -194,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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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