Get the db admin green agent working

2024-02-27 13:30:16 +00:00
parent c54f82fb1b
commit af8ca82fcb
16 changed files with 386 additions and 87 deletions
--- a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py
+++ b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py
@@ -0,0 +1,84 @@
+from primaite.game.agent.actions import ActionManager
+from primaite.game.agent.observations import ICSObservation, ObservationManager
+from primaite.game.agent.rewards import RewardFunction
+from primaite.game.agent.scripted_agents import ProbabilisticAgent
+
+
+def test_probabilistic_agent():
+    """
+    Check that the probabilistic agent selects actions with approximately the right probabilities.
+
+    Using a binomial probability calculator (https://www.wolframalpha.com/input/?i=binomial+distribution+calculator),
+    we can generate some lower and upper bounds of how many times we expect the agent to take each action. These values
+    were chosen to guarantee a less than 1 in a million chance of the test failing due to unlucky random number
+    generation.
+    """
+    N_TRIALS = 10_000
+    P_DO_NOTHING = 0.1
+    P_NODE_APPLICATION_EXECUTE = 0.3
+    P_NODE_FILE_DELETE = 0.6
+    MIN_DO_NOTHING = 850
+    MAX_DO_NOTHING = 1150
+    MIN_NODE_APPLICATION_EXECUTE = 2800
+    MAX_NODE_APPLICATION_EXECUTE = 3200
+    MIN_NODE_FILE_DELETE = 5750
+    MAX_NODE_FILE_DELETE = 6250
+
+    action_space = ActionManager(
+        actions=[
+            {"type": "DONOTHING"},
+            {"type": "NODE_APPLICATION_EXECUTE"},
+            {"type": "NODE_FILE_DELETE"},
+        ],
+        nodes=[
+            {
+                "node_name": "client_1",
+                "applications": [{"application_name": "WebBrowser"}],
+                "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
+            },
+        ],
+        max_folders_per_node=2,
+        max_files_per_folder=2,
+        max_services_per_node=2,
+        max_applications_per_node=2,
+        max_nics_per_node=2,
+        max_acl_rules=10,
+        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}},
+        },
+    )
+    observation_space = ObservationManager(ICSObservation())
+    reward_function = RewardFunction()
+
+    pa = ProbabilisticAgent(
+        agent_name="test_agent",
+        action_space=action_space,
+        observation_space=observation_space,
+        reward_function=reward_function,
+        settings={
+            "action_probabilities": {0: P_DO_NOTHING, 1: P_NODE_APPLICATION_EXECUTE, 2: P_NODE_FILE_DELETE},
+            "random_seed": 120,
+        },
+    )
+
+    do_nothing_count = 0
+    node_application_execute_count = 0
+    node_file_delete_count = 0
+    for _ in range(N_TRIALS):
+        a = pa.get_action(0, timestep=0)
+        if a == ("DONOTHING", {}):
+            do_nothing_count += 1
+        elif a == ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}):
+            node_application_execute_count += 1
+        elif a == ("NODE_FILE_DELETE", {"node_id": 0, "folder_id": 0, "file_id": 0}):
+            node_file_delete_count += 1
+        else:
+            raise AssertionError("Probabilistic agent produced an unexpected action.")
+
+    assert MIN_DO_NOTHING < do_nothing_count < MAX_DO_NOTHING
+    assert MIN_NODE_APPLICATION_EXECUTE < node_application_execute_count < MAX_NODE_APPLICATION_EXECUTE
+    assert MIN_NODE_FILE_DELETE < node_file_delete_count < MAX_NODE_FILE_DELETE