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

Fix software reset issues

Browse Source
This commit is contained in:
Marek Wolan
2023-11-29 13:18:38 +00:00
parent bf73cc2eb7
commit 05d62a956d
15 changed files with 119 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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: