From 80b1988ab93306d620a47ee34478ece07f41e61b Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Mon, 20 May 2024 13:10:21 +0100 Subject: [PATCH] #2606: add agent actions + test --- src/primaite/game/agent/actions.py | 42 ++++++++++- .../simulator/file_system/file_system.py | 36 +++++++--- tests/conftest.py | 3 + .../game_layer/test_actions.py | 72 +++++++++++++++++++ 4 files changed, 141 insertions(+), 12 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 430b5b1d..f6220465 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -313,13 +313,35 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction): self.verb: str = "restore" -class NodeFolderFileCreateAction(NodeFolderAbstractAction): +class NodeFileCreateAction(AbstractAction): """Action which creates a new file in a given folder.""" def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) self.verb: str = "create" + def form_request(self, node_id: int, folder_name: str, file_name: str) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None or folder_name is None or file_name is None: + return ["do_nothing"] + return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name] + + +class NodeFolderCreateAction(AbstractAction): + """Action which creates a new folder.""" + + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) + self.verb: str = "create" + + def form_request(self, node_id: int, folder_name: str) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None or folder_name is None: + return ["do_nothing"] + return ["network", "node", node_name, "file_system", "create", "folder", folder_name] + class NodeFileAbstractAction(AbstractAction): """Abstract base class for file actions. @@ -401,6 +423,21 @@ class NodeFileCorruptAction(NodeFileAbstractAction): self.verb: str = "corrupt" +class NodeFileAccessAction(AbstractAction): + """Action which increases a file's access count.""" + + def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: + super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) + self.verb: str = "create" + + def form_request(self, node_id: int, folder_name: str, file_name: str) -> List[str]: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + if node_name is None or folder_name is None or file_name is None: + return ["do_nothing"] + return ["network", "node", node_name, "file_system", "access", folder_name, file_name] + + class NodeAbstractAction(AbstractAction): """ Abstract base class for node actions. @@ -854,11 +891,14 @@ class ActionManager: "NODE_APPLICATION_INSTALL": NodeApplicationInstallAction, "NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction, "NODE_FILE_SCAN": NodeFileScanAction, + "NODE_FILE_CREATE": NodeFileCreateAction, "NODE_FILE_CHECKHASH": NodeFileCheckhashAction, "NODE_FILE_DELETE": NodeFileDeleteAction, "NODE_FILE_REPAIR": NodeFileRepairAction, "NODE_FILE_RESTORE": NodeFileRestoreAction, "NODE_FILE_CORRUPT": NodeFileCorruptAction, + "NODE_FILE_ACCESS": NodeFileAccessAction, + "NODE_FOLDER_CREATE": NodeFolderCreateAction, "NODE_FOLDER_SCAN": NodeFolderScanAction, "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, "NODE_FOLDER_REPAIR": NodeFolderRepairAction, diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 8240cba3..7c424d22 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, List, Optional from prettytable import MARKDOWN, PrettyTable @@ -64,21 +64,35 @@ class FileSystem(SimComponent): ) self._create_manager = RequestManager() + + def _create_file_action(request: List[Any], context: Any) -> RequestResponse: + file = self.create_file(folder_name=request[0], file_name=request[1]) + if not file: + return RequestResponse.from_bool(False) + return RequestResponse( + status="success", + data={ + "file_name": file.name, + "folder_name": file.folder_name, + "file_type": file.file_type, + "file_size": file.size, + }, + ) + self._create_manager.add_request( name="file", - request_type=RequestType( - func=lambda request, context: RequestResponse.from_bool( - bool(self.create_file(folder_name=request[0], file_name=request[1])) - ) - ), + request_type=RequestType(func=_create_file_action), ) + + def _create_folder_action(request: List[Any], context: Any) -> RequestResponse: + folder = self.create_folder(folder_name=request[0]) + if not folder: + return RequestResponse.from_bool(False) + return RequestResponse(status="success", data={"folder_name": folder.name}) + self._create_manager.add_request( name="folder", - request_type=RequestType( - func=lambda request, context: RequestResponse.from_bool( - bool(self.create_folder(folder_name=request[0])) - ) - ), + request_type=RequestType(func=_create_folder_action), ) rm.add_request( name="create", diff --git a/tests/conftest.py b/tests/conftest.py index 83778748..a20822f3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -431,12 +431,15 @@ def game_and_agent(): {"type": "NODE_APPLICATION_FIX"}, {"type": "NODE_APPLICATION_INSTALL"}, {"type": "NODE_APPLICATION_REMOVE"}, + {"type": "NODE_FILE_CREATE"}, {"type": "NODE_FILE_SCAN"}, {"type": "NODE_FILE_CHECKHASH"}, {"type": "NODE_FILE_DELETE"}, {"type": "NODE_FILE_REPAIR"}, {"type": "NODE_FILE_RESTORE"}, {"type": "NODE_FILE_CORRUPT"}, + {"type": "NODE_FILE_ACCESS"}, + {"type": "NODE_FOLDER_CREATE"}, {"type": "NODE_FOLDER_SCAN"}, {"type": "NODE_FOLDER_CHECKHASH"}, {"type": "NODE_FOLDER_REPAIR"}, diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index edaf5d8d..43987c38 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -329,6 +329,78 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA assert file.deleted +def test_node_file_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that a file is created.""" + game, agent = game_and_agent + + client_1 = game.simulation.network.get_node_by_hostname("client_1") # + + action = ( + "NODE_FILE_CREATE", + { + "node_id": 0, + "folder_name": "test", + "file_name": "file.txt", + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_file(folder_name="test", file_name="file.txt") + + +def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the file access increments.""" + game, agent = game_and_agent + + client_1 = game.simulation.network.get_node_by_hostname("client_1") # + + action = ( + "NODE_FILE_CREATE", + { + "node_id": 0, + "folder_name": "test", + "file_name": "file.txt", + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 0 + + action = ( + "NODE_FILE_ACCESS", + { + "node_id": 0, + "folder_name": "test", + "file_name": "file.txt", + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 1 + + +def test_node_folder_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that a folder is created.""" + game, agent = game_and_agent + + client_1 = game.simulation.network.get_node_by_hostname("client_1") # + + action = ( + "NODE_FOLDER_CREATE", + { + "node_id": 0, + "folder_name": "test", + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_folder(folder_name="test") + + def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): """Test that the NetworkPortDisableAction can form a request and that it is accepted by the simulation.""" game, agent = game_and_agent