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

#1859 - As disccused

Browse Source
This commit is contained in:
Chris McCarthy
2023-11-27 17:05:12 +00:00
parent c3bb78c5a3
commit ae5046b8fb
7 changed files with 83 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/primaite/game/agent/actions.py
View File

@@ -634,7 +634,6 @@ class ActionManager:
:type act_map: Optional[Dict[int, Dict]]
"""
self.game: "PrimaiteGame" = game
self.sim: Simulation = self.game.simulation
self.node_uuids: List[str] = node_uuids
self.application_uuids: List[List[str]] = application_uuids
self.protocols: List[str] = protocols
@@ -646,7 +645,7 @@ class ActionManager:
else:
self.ip_address_list = []
for node_uuid in self.node_uuids:
node_obj = self.sim.network.nodes[node_uuid]
node_obj = self.game.simulation.network.nodes[node_uuid]
nics = node_obj.nics
for nic_uuid, nic_obj in nics.items():
self.ip_address_list.append(nic_obj.ip_address)
@@ -770,7 +769,7 @@ class ActionManager:
:rtype: Optional[str]
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.sim.network.nodes[node_uuid]
node = self.game.simulation.network.nodes[node_uuid]
folder_uuids = list(node.file_system.folders.keys())
return folder_uuids[folder_idx] if len(folder_uuids) > folder_idx else None
@@ -788,7 +787,7 @@ class ActionManager:
:rtype: Optional[str]
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.sim.network.nodes[node_uuid]
node = self.game.simulation.network.nodes[node_uuid]
folder_uuids = list(node.file_system.folders.keys())
if len(folder_uuids) <= folder_idx:
return None
@@ -807,7 +806,7 @@ class ActionManager:
:rtype: Optional[str]
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.sim.network.nodes[node_uuid]
node = self.game.simulation.network.nodes[node_uuid]
service_uuids = list(node.services.keys())
return service_uuids[service_idx] if len(service_uuids) > service_idx else None
@@ -867,7 +866,7 @@ class ActionManager:
:rtype: str
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node_obj = self.sim.network.nodes[node_uuid]
node_obj = self.game.simulation.network.nodes[node_uuid]
nics = list(node_obj.nics.keys())
if len(nics) <= nic_idx:
return None

2
src/primaite/game/agent/observations.py
View File

@@ -162,7 +162,7 @@ class ServiceObservation(AbstractObservation):
:return: Constructed service observation
:rtype: ServiceObservation
"""
return cls(where=parent_where + ["services", game.ref_map_services[config["service_ref"]].uuid])
return cls(where=parent_where + ["services", game.ref_map_services[config["service_ref"]]])
class LinkObservation(AbstractObservation):

3
src/primaite/game/agent/rewards.py
View File

@@ -25,6 +25,7 @@ the structure:
service_ref: web_server_database_client
```
"""
import json
from abc import abstractmethod
from typing import Dict, List, Tuple, Type, TYPE_CHECKING
@@ -213,7 +214,7 @@ class WebServer404Penalty(AbstractReward):
_LOGGER.warn(msg)
return DummyReward() # TODO: should we error out with incorrect inputs? Probably!
node_uuid = game.ref_map_nodes[node_ref]
service_uuid = game.ref_map_services[service_ref].uuid
service_uuid = game.ref_map_services[service_ref]
if not (node_uuid and service_uuid):
msg = (
f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not"

56
src/primaite/game/game.py
View File

@@ -59,8 +59,9 @@ class PrimaiteGame:
"""Initialise a PrimaiteGame object."""
self.simulation: Simulation = Simulation()
"""Simulation object with which the agents will interact."""
print(f"Hello, welcome to PrimaiteGame. This is the ID of the ORIGINAL simulation {id(self.simulation)}")
self._simulation_initial_state = deepcopy(self.simulation)
self._simulation_initial_state = None
"""The Simulation original state (deepcopy of the original Simulation)."""
self.agents: List[AbstractAgent] = []
@@ -78,16 +79,16 @@ class PrimaiteGame:
self.options: PrimaiteGameOptions
"""Special options that apply for the entire game."""
self.ref_map_nodes: Dict[str, Node] = {}
self.ref_map_nodes: Dict[str, str] = {}
"""Mapping from unique node reference name to node object. Used when parsing config files."""
self.ref_map_services: Dict[str, Service] = {}
self.ref_map_services: Dict[str, str] = {}
"""Mapping from human-readable service reference to service object. Used for parsing config files."""
self.ref_map_applications: Dict[str, Application] = {}
self.ref_map_applications: Dict[str, str] = {}
"""Mapping from human-readable application reference to application object. Used for parsing config files."""
self.ref_map_links: Dict[str, Link] = {}
self.ref_map_links: Dict[str, str] = {}
"""Mapping from human-readable link reference to link object. Used when parsing config files."""
def step(self):
@@ -161,6 +162,33 @@ class PrimaiteGame:
self.step_counter = 0
_LOGGER.debug(f"Resetting primaite game, episode = {self.episode_counter}")
self.simulation = deepcopy(self._simulation_initial_state)
self._reset_components_for_episode()
print("Reset")
def _reset_components_for_episode(self):
print("Performing full reset for episode")
for node in self.simulation.network.nodes.values():
print(f"Resetting Node: {node.hostname}")
node.reset_component_for_episode(self.episode_counter)
# reset Node NIC
# Reset Node Services
# Reset Node Applications
print(f"Resetting Software...")
for application in node.software_manager.software.values():
print(f"Resetting {application.name}")
if isinstance(application, WebBrowser):
application.do_this()
# Reset Node FileSystem
# Reset Node FileSystemFolder's
# Reset Node FileSystemFile's
# Reset Router
# Reset Links
def close(self) -> None:
"""Close the game, this will close the simulation."""
@@ -190,10 +218,6 @@ class PrimaiteGame:
sim = game.simulation
net = sim.network
game.ref_map_nodes: Dict[str, Node] = {}
game.ref_map_services: Dict[str, Service] = {}
game.ref_map_links: Dict[str, Link] = {}
nodes_cfg = cfg["simulation"]["network"]["nodes"]
links_cfg = cfg["simulation"]["network"]["links"]
for node_cfg in nodes_cfg:
@@ -269,7 +293,7 @@ class PrimaiteGame:
print(f"installing {service_type} on node {new_node.hostname}")
new_node.software_manager.install(service_types_mapping[service_type])
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service
game.ref_map_services[service_ref] = new_service.uuid
else:
print(f"service type not found {service_type}")
# service-dependent options
@@ -303,7 +327,7 @@ class PrimaiteGame:
if application_type in application_types_mapping:
new_node.software_manager.install(application_types_mapping[application_type])
new_application = new_node.software_manager.software[application_type]
game.ref_map_applications[application_ref] = new_application
game.ref_map_applications[application_ref] = new_application.uuid
else:
print(f"application type not found {application_type}")
@@ -326,11 +350,7 @@ class PrimaiteGame:
net.add_node(new_node)
new_node.power_on()
game.ref_map_nodes[
node_ref
] = (
new_node.uuid
) # TODO: fix inconsistency with service and link. Node gets added by uuid, but service by object
game.ref_map_nodes[node_ref] = new_node.uuid
# 2. create links between nodes
for link_cfg in links_cfg:
@@ -375,7 +395,7 @@ class PrimaiteGame:
for application_option in action_node_option["applications"]:
# TODO: fix inconsistency with node uuids and application uuids. The node object get added to
# node_uuid, whereas here the application gets added by uuid.
application_uuid = game.ref_map_applications[application_option["application_ref"]].uuid
application_uuid = game.ref_map_applications[application_option["application_ref"]]
node_application_uuids.append(application_uuid)
action_space_cfg["options"]["application_uuids"].append(node_application_uuids)
@@ -433,5 +453,7 @@ class PrimaiteGame:
print("agent type not found")
game._simulation_initial_state = deepcopy(game.simulation) # noqa
web_server = game.simulation.network.get_node_by_hostname("web_server").software_manager.software["WebServer"]
print(f"And this is the ID of the original WebServer {id(web_server)}")
return game

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

@@ -1005,6 +1005,9 @@ class Node(SimComponent):
return rm
def reset_component_for_episode(self, episode: int):
self._init_request_manager()
def _install_system_software(self):
"""Install System Software - software that is usually provided with the OS."""
pass

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

@@ -43,6 +43,13 @@ class WebBrowser(Application):
return rm
def do_this(self):
self._init_request_manager()
print(f"Resetting WebBrowser for episode")
def reset_component_for_episode(self, episode: int):
pass
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of the WebBrowser.

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

@@ -17,7 +17,16 @@ 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
@property
def last_response_status_code(self) -> HttpStatusCode:
return self._last_response_status_code
@last_response_status_code.setter
def last_response_status_code(self, val: Any):
print(f"val: {val}, type: {type(val)}")
self._last_response_status_code = val
def describe_state(self) -> Dict:
"""
@@ -29,10 +38,17 @@ class WebServer(Service):
:rtype: Dict
"""
state = super().describe_state()
state["last_response_status_code"] = (
self.last_response_status_code.value if isinstance(self.last_response_status_code, HttpStatusCode) else None
)
print(
f""
f"Printing state from Webserver describe func: "
f"val={state['last_response_status_code']}, "
f"type={type(state['last_response_status_code'])}, "
f"Service obj ID={id(self)}"
)
return state
def __init__(self, **kwargs):
@@ -85,7 +101,14 @@ class WebServer(Service):
# return true if response is OK
self.last_response_status_code = response.status_code
print(self.last_response_status_code)
print(
f""
f"Printing state from Webserver http request func: "
f"val={self.last_response_status_code}, "
f"type={type(self.last_response_status_code)}, "
f"Service obj ID={id(self)}"
)
return response.status_code == HttpStatusCode.OK
def _handle_get_request(self, payload: HttpRequestPacket) -> HttpResponsePacket: