From d3c52d0d7296ddd18f4f6fb7d84856a169577849 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 19 Nov 2024 18:58:51 +0000 Subject: [PATCH] #2912 - Remove some debugging print statements and apply pre-commit lint changes --- src/primaite/game/agent/actions/abstract.py | 3 +-- src/primaite/game/agent/actions/acl.py | 5 ++--- src/primaite/game/agent/actions/config.py | 15 ++++++++++----- src/primaite/game/agent/actions/host_nic.py | 1 + src/primaite/game/agent/actions/manager.py | 3 +-- src/primaite/game/agent/actions/node.py | 2 +- .../integration_tests/game_layer/test_actions.py | 4 ++-- 7 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/primaite/game/agent/actions/abstract.py b/src/primaite/game/agent/actions/abstract.py index 5c0594fd..cd14ef6d 100644 --- a/src/primaite/game/agent/actions/abstract.py +++ b/src/primaite/game/agent/actions/abstract.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, ConfigDict from primaite.interface.request import RequestFormat + class AbstractAction(BaseModel): """Base class for actions.""" @@ -37,6 +38,4 @@ class AbstractAction(BaseModel): """Create an action component from a config dictionary.""" if not config.get("type"): config.update({"type": cls.__name__}) - print("oooh") - print(config) return cls(config=cls.ConfigSchema(**config)) diff --git a/src/primaite/game/agent/actions/acl.py b/src/primaite/game/agent/actions/acl.py index 11269a7e..7ab49732 100644 --- a/src/primaite/game/agent/actions/acl.py +++ b/src/primaite/game/agent/actions/acl.py @@ -127,8 +127,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r if config.src_ip == 0: return ["do_nothing"] # invalid formulation if config.src_port == 0: - return ["do_nothing"] # invalid configuration. - + return ["do_nothing"] # invalid configuration. return [ "network", @@ -153,7 +152,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r class FirewallACLRemoveRuleAction(AbstractAction, identifier="firewall_acl_remove_rule"): """Action which removes a rule from a firewall port's ACL.""" - config:"FirewallACLRemoveRuleAction.ConfigSchema" + config: "FirewallACLRemoveRuleAction.ConfigSchema" class ConfigSchema(AbstractAction.ConfigSchema): """Configuration schema for FirewallACLRemoveRuleAction.""" diff --git a/src/primaite/game/agent/actions/config.py b/src/primaite/game/agent/actions/config.py index da9f77e6..7c72e57d 100644 --- a/src/primaite/game/agent/actions/config.py +++ b/src/primaite/game/agent/actions/config.py @@ -38,7 +38,15 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="c2_server_rans """Return the action formatted as a request that can be ingested by the simulation.""" if config.node_name is None: return ["do_nothing"] - return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", config.model_config] + return [ + "network", + "node", + config.node_name, + "application", + "RansomwareScript", + "configure", + config.model_config, + ] class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): @@ -207,10 +215,7 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi exfiltration_folder_name: Optional[str] @classmethod - def form_request( - cls, - config: ConfigSchema - ) -> RequestFormat: + def form_request(cls, config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" if config.node_name is None: return ["do_nothing"] diff --git a/src/primaite/game/agent/actions/host_nic.py b/src/primaite/game/agent/actions/host_nic.py index 6df241bc..e2adf7d7 100644 --- a/src/primaite/game/agent/actions/host_nic.py +++ b/src/primaite/game/agent/actions/host_nic.py @@ -1,5 +1,6 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from typing import ClassVar + from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py index 9ef94069..a413f6dc 100644 --- a/src/primaite/game/agent/actions/manager.py +++ b/src/primaite/game/agent/actions/manager.py @@ -31,7 +31,6 @@ class DoNothingAction(AbstractAction, identifier="do_nothing"): class ConfigSchema(AbstractAction.ConfigSchema): """Configuration Schema for DoNothingAction.""" - # type: Literal["do_nothing"] = "do_nothing" type: str = "do_nothing" @classmethod @@ -133,7 +132,7 @@ class ActionManager: def space(self) -> spaces.Space: """Return the gymnasium action space for this agent.""" return spaces.Discrete(len(self.action_map)) - + @classmethod def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager": """ diff --git a/src/primaite/game/agent/actions/node.py b/src/primaite/game/agent/actions/node.py index 3c70d495..4ecc1393 100644 --- a/src/primaite/game/agent/actions/node.py +++ b/src/primaite/game/agent/actions/node.py @@ -90,7 +90,7 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti """Base Configuration Schema for NodeNMAP actions.""" target_ip_address: Union[str, List[str]] - show: bool = False + show: bool = False node_name: str @classmethod diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index a31f325a..a21ad34f 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -56,7 +56,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.UNUSED # 2: Scan and check that the visible state is now correct - action = ("node_service_scan", {"type":"node_service_scan" ,"node_name": "server_1", "service_name": "DNSServer"}) + action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.GOOD @@ -67,7 +67,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.GOOD # 4: Scan and check that the visible state is now correct - action = ("node_service_scan", {"type":"node_service_scan", "node_name": "server_1", "service_name": "DNSServer"}) + action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.COMPROMISED