#2606: add agent actions + test

This commit is contained in:
Czar Echavez
2024-05-20 13:10:21 +01:00
parent 406d4e889b
commit 80b1988ab9
4 changed files with 141 additions and 12 deletions

View File

@@ -313,13 +313,35 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction):
self.verb: str = "restore" self.verb: str = "restore"
class NodeFolderFileCreateAction(NodeFolderAbstractAction): class NodeFileCreateAction(AbstractAction):
"""Action which creates a new file in a given folder.""" """Action which creates a new file in a given folder."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: 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) super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
self.verb: str = "create" 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): class NodeFileAbstractAction(AbstractAction):
"""Abstract base class for file actions. """Abstract base class for file actions.
@@ -401,6 +423,21 @@ class NodeFileCorruptAction(NodeFileAbstractAction):
self.verb: str = "corrupt" 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): class NodeAbstractAction(AbstractAction):
""" """
Abstract base class for node actions. Abstract base class for node actions.
@@ -854,11 +891,14 @@ class ActionManager:
"NODE_APPLICATION_INSTALL": NodeApplicationInstallAction, "NODE_APPLICATION_INSTALL": NodeApplicationInstallAction,
"NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction, "NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction,
"NODE_FILE_SCAN": NodeFileScanAction, "NODE_FILE_SCAN": NodeFileScanAction,
"NODE_FILE_CREATE": NodeFileCreateAction,
"NODE_FILE_CHECKHASH": NodeFileCheckhashAction, "NODE_FILE_CHECKHASH": NodeFileCheckhashAction,
"NODE_FILE_DELETE": NodeFileDeleteAction, "NODE_FILE_DELETE": NodeFileDeleteAction,
"NODE_FILE_REPAIR": NodeFileRepairAction, "NODE_FILE_REPAIR": NodeFileRepairAction,
"NODE_FILE_RESTORE": NodeFileRestoreAction, "NODE_FILE_RESTORE": NodeFileRestoreAction,
"NODE_FILE_CORRUPT": NodeFileCorruptAction, "NODE_FILE_CORRUPT": NodeFileCorruptAction,
"NODE_FILE_ACCESS": NodeFileAccessAction,
"NODE_FOLDER_CREATE": NodeFolderCreateAction,
"NODE_FOLDER_SCAN": NodeFolderScanAction, "NODE_FOLDER_SCAN": NodeFolderScanAction,
"NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction, "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction,
"NODE_FOLDER_REPAIR": NodeFolderRepairAction, "NODE_FOLDER_REPAIR": NodeFolderRepairAction,

View File

@@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
from pathlib import Path from pathlib import Path
from typing import Dict, Optional from typing import Any, Dict, List, Optional
from prettytable import MARKDOWN, PrettyTable from prettytable import MARKDOWN, PrettyTable
@@ -64,21 +64,35 @@ class FileSystem(SimComponent):
) )
self._create_manager = RequestManager() 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( self._create_manager.add_request(
name="file", name="file",
request_type=RequestType( request_type=RequestType(func=_create_file_action),
func=lambda request, context: RequestResponse.from_bool(
bool(self.create_file(folder_name=request[0], file_name=request[1]))
)
),
) )
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( self._create_manager.add_request(
name="folder", name="folder",
request_type=RequestType( request_type=RequestType(func=_create_folder_action),
func=lambda request, context: RequestResponse.from_bool(
bool(self.create_folder(folder_name=request[0]))
)
),
) )
rm.add_request( rm.add_request(
name="create", name="create",

View File

@@ -431,12 +431,15 @@ def game_and_agent():
{"type": "NODE_APPLICATION_FIX"}, {"type": "NODE_APPLICATION_FIX"},
{"type": "NODE_APPLICATION_INSTALL"}, {"type": "NODE_APPLICATION_INSTALL"},
{"type": "NODE_APPLICATION_REMOVE"}, {"type": "NODE_APPLICATION_REMOVE"},
{"type": "NODE_FILE_CREATE"},
{"type": "NODE_FILE_SCAN"}, {"type": "NODE_FILE_SCAN"},
{"type": "NODE_FILE_CHECKHASH"}, {"type": "NODE_FILE_CHECKHASH"},
{"type": "NODE_FILE_DELETE"}, {"type": "NODE_FILE_DELETE"},
{"type": "NODE_FILE_REPAIR"}, {"type": "NODE_FILE_REPAIR"},
{"type": "NODE_FILE_RESTORE"}, {"type": "NODE_FILE_RESTORE"},
{"type": "NODE_FILE_CORRUPT"}, {"type": "NODE_FILE_CORRUPT"},
{"type": "NODE_FILE_ACCESS"},
{"type": "NODE_FOLDER_CREATE"},
{"type": "NODE_FOLDER_SCAN"}, {"type": "NODE_FOLDER_SCAN"},
{"type": "NODE_FOLDER_CHECKHASH"}, {"type": "NODE_FOLDER_CHECKHASH"},
{"type": "NODE_FOLDER_REPAIR"}, {"type": "NODE_FOLDER_REPAIR"},

View File

@@ -329,6 +329,78 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
assert file.deleted 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]): 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.""" """Test that the NetworkPortDisableAction can form a request and that it is accepted by the simulation."""
game, agent = game_and_agent game, agent = game_and_agent