Merge agents and actions branches + fix import / subclass errors

2025-01-14 11:34:01 +00:00
parent 9f5e16dd85 1ac562ebc9
commit fb5242bd78
57 changed files with 2297 additions and 2381 deletions
--- a/tests/unit_tests/_primaite/_game/_agent/test_actions.py
+++ b/tests/unit_tests/_primaite/_game/_agent/test_actions.py
@@ -3,9 +3,11 @@ from unittest.mock import Mock

 import pytest

-from primaite.game.agent.actions import (
+from primaite.game.agent.actions import (  # DoNothingAction,; NodeServiceDisableAction,; NodeServiceEnableAction,; NodeServicePauseAction,; NodeServiceRestartAction,; NodeServiceResumeAction,; NodeServiceScanAction,; NodeServiceStartAction,; NodeServiceStopAction,
    ActionManager,
-    DoNothingAction,
+)
+from primaite.game.agent.actions.manager import DoNothingAction
+from primaite.game.agent.actions.service import (
    NodeServiceDisableAction,
    NodeServiceEnableAction,
    NodeServicePauseAction,
@@ -18,7 +20,7 @@ from primaite.game.agent.actions import (


 def test_do_nothing_action_form_request():
-    """Test that the DoNothingAction can form a request and that it is correct."""
+    """Test that the do_nothingAction can form a request and that it is correct."""
    manager = Mock()

    action = DoNothingAction(manager=manager)
--- a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py
+++ b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py
@@ -28,9 +28,9 @@ def test_probabilistic_agent():

    action_space_cfg = {
        "action_list": [
-            {"type": "DONOTHING"},
-            {"type": "NODE_APPLICATION_EXECUTE"},
-            {"type": "NODE_FILE_DELETE"},
+            {"type": "do_nothing"},
+            {"type": "node_application_execute"},
+            {"type": "node_file_delete"},
        ],
        "nodes": [
            {
@@ -48,9 +48,9 @@ def test_probabilistic_agent():
        "protocols": ["TCP", "UDP", "ICMP"],
        "ports": ["HTTP", "DNS", "ARP"],
        "act_map": {
-            0: {"action": "DONOTHING", "options": {}},
-            1: {"action": "NODE_APPLICATION_EXECUTE", "options": {"node_id": 0, "application_id": 0}},
-            2: {"action": "NODE_FILE_DELETE", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
+            0: {"action": "do_nothing", "options": {}},
+            1: {"action": "node_application_execute", "options": {"node_id": 0, "application_id": 0}},
+            2: {"action": "node_file_delete", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
        },
        "options": {},
    }
@@ -80,11 +80,11 @@ def test_probabilistic_agent():
    node_file_delete_count = 0
    for _ in range(N_TRIALS):
        a = pa.get_action(0)
-        if a == ("DONOTHING", {}):
+        if a == ("do_nothing", {}):
            do_nothing_count += 1
-        elif a == ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}):
+        elif a == ("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"}):
            node_application_execute_count += 1
-        elif a == ("NODE_FILE_DELETE", {"node_id": 0, "folder_id": 0, "file_id": 0}):
+        elif a == ("node_file_delete", {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}):
            node_file_delete_count += 1
        else:
            raise AssertionError("Probabilistic agent produced an unexpected action.")
--- a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py
+++ b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py
@@ -11,7 +11,12 @@ from primaite.interface.request import RequestResponse

 class TestWebServer404PenaltySticky:
    def test_non_sticky(self):
-        reward = WebServer404Penalty("computer", "WebService", sticky=False)
+        schema = WebServer404Penalty.ConfigSchema(
+            node_hostname="computer",
+            service_name="WebService",
+            sticky=False,
+        )
+        reward = WebServer404Penalty(config=schema)

        # no response codes yet, reward is 0
        codes = []
@@ -38,7 +43,12 @@ class TestWebServer404PenaltySticky:
        assert reward.calculate(state, last_action_response) == -1.0

    def test_sticky(self):
-        reward = WebServer404Penalty("computer", "WebService", sticky=True)
+        schema = WebServer404Penalty.ConfigSchema(
+            node_hostname="computer",
+            service_name="WebService",
+            sticky=True,
+        )
+        reward = WebServer404Penalty(config=schema)

        # no response codes yet, reward is 0
        codes = []
@@ -67,10 +77,11 @@ class TestWebServer404PenaltySticky:

 class TestWebpageUnavailabilitySticky:
    def test_non_sticky(self):
-        reward = WebpageUnavailablePenalty("computer", sticky=False)
+        schema = WebpageUnavailablePenalty.ConfigSchema(node_hostname="computer", sticky=False)
+        reward = WebpageUnavailablePenalty(config=schema)

        # no response codes yet, reward is 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "do_nothing", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        browser_history = []
        state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -93,7 +104,7 @@ class TestWebpageUnavailabilitySticky:

        # THE IMPORTANT BIT
        # agent did nothing, because reward is not sticky, it goes back to 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        browser_history = []
        state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -127,10 +138,11 @@ class TestWebpageUnavailabilitySticky:
        assert reward.calculate(state, last_action_response) == -1.0

    def test_sticky(self):
-        reward = WebpageUnavailablePenalty("computer", sticky=True)
+        schema = WebpageUnavailablePenalty.ConfigSchema(node_hostname="computer", sticky=True)
+        reward = WebpageUnavailablePenalty(config=schema)

        # no response codes yet, reward is 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        browser_history = []
        state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -153,7 +165,7 @@ class TestWebpageUnavailabilitySticky:

        # THE IMPORTANT BIT
        # agent did nothing, because reward is sticky, it stays at 1.0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
        last_action_response = AgentHistoryItem(
@@ -188,10 +200,14 @@ class TestWebpageUnavailabilitySticky:

 class TestGreenAdminDatabaseUnreachableSticky:
    def test_non_sticky(self):
-        reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=False)
+        schema = GreenAdminDatabaseUnreachablePenalty.ConfigSchema(
+            node_hostname="computer",
+            sticky=False,
+        )
+        reward = GreenAdminDatabaseUnreachablePenalty(config=schema)

        # no response codes yet, reward is 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
        last_action_response = AgentHistoryItem(
@@ -212,9 +228,8 @@ class TestGreenAdminDatabaseUnreachableSticky:

        # THE IMPORTANT BIT
        # agent did nothing, because reward is not sticky, it goes back to 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
-        browser_history = []
        state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
        last_action_response = AgentHistoryItem(
            timestep=0, action=action, parameters=params, request=request, response=response
@@ -244,10 +259,14 @@ class TestGreenAdminDatabaseUnreachableSticky:
        assert reward.calculate(state, last_action_response) == -1.0

    def test_sticky(self):
-        reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=True)
+        schema = GreenAdminDatabaseUnreachablePenalty.ConfigSchema(
+            node_hostname="computer",
+            sticky=True,
+        )
+        reward = GreenAdminDatabaseUnreachablePenalty(config=schema)

        # no response codes yet, reward is 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
        last_action_response = AgentHistoryItem(
@@ -268,7 +287,7 @@ class TestGreenAdminDatabaseUnreachableSticky:

        # THE IMPORTANT BIT
        # agent did nothing, because reward is not sticky, it goes back to 0
-        action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
+        action, params, request = "DO_NOTHING", {}, ["do_nothing"]
        response = RequestResponse(status="success", data={})
        state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
        last_action_response = AgentHistoryItem(