Apply suggestions from PR review.

2024-03-04 09:58:57 +00:00
parent a6031d568d
commit d1480e4477
16 changed files with 25 additions and 19 deletions
--- a/docs/source/configuration/agents.rst
+++ b/docs/source/configuration/agents.rst
@@ -19,7 +19,7 @@ Agents can be scripted (deterministic and stochastic), or controlled by a reinfo
            ...
        - ref: green_agent_example
            team: GREEN
-            type: probabilistic_agent
+            type: ProbabilisticAgent
            observation_space:
                type: UC2GreenObservation
            action_space:
@@ -57,11 +57,11 @@ Specifies if the agent is malicious (``RED``), benign (``GREEN``), or defensive

 ``type``
 --------
-Specifies which class should be used for the agent. ``ProxyAgent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``RedDatabaseCorruptingAgent`` and ``probabilistic_agent`` generate their own behaviour.
+Specifies which class should be used for the agent. ``ProxyAgent`` is used for agents that receive instructions from an RL algorithm. Scripted agents like ``RedDatabaseCorruptingAgent`` and ``ProbabilisticAgent`` generate their own behaviour.

 Available agent types:

- ``probabilistic_agent``
+- ``ProbabilisticAgent``
 - ``ProxyAgent``
 - ``RedDatabaseCorruptingAgent``

--- a/docs/source/simulation_components/system/services/database_service.rst
+++ b/docs/source/simulation_components/system/services/database_service.rst
@@ -25,6 +25,13 @@ Usage
 - Clients connect, execute queries, and disconnect.
 - Service runs on TCP port 5432 by default.

+**Supported queries:**
+
+* ``SELECT``: As long as the database file is in a ``GOOD`` health state, the db service will respond with a 200 status code.
+* ``DELETE``: This query represents an attack, it will cause the database file to enter a ``COMPROMISED`` state, and return a status code 200.
+* ``INSERT``: If the database service is in a healthy state, this will return a 200 status, if it's not in a healthy state it will return 404.
+* ``SELECT * FROM pg_stat_activity``: This query represents something an admin would send to check the status of the database. If the database service is in a healthy state, it returns a 200 status code, otherwise a 401 status code.
+
 Implementation
 ==============

--- a/src/primaite/config/_package_data/example_config.yaml
+++ b/src/primaite/config/_package_data/example_config.yaml
@@ -33,7 +33,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    agent_settings:
      action_probabilities:
        0: 0.3
@@ -76,7 +76,7 @@ agents:

  - ref: client_1_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    agent_settings:
      action_probabilities:
        0: 0.3
--- a/src/primaite/config/_package_data/example_config_2_rl_agents.yaml
+++ b/src/primaite/config/_package_data/example_config_2_rl_agents.yaml
@@ -35,7 +35,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
@@ -64,7 +64,7 @@ agents:

  - ref: client_1_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/src/primaite/game/game.py
+++ b/src/primaite/game/game.py
@@ -414,7 +414,7 @@ class PrimaiteGame:
            reward_function = RewardFunction.from_config(reward_function_cfg)

            # CREATE AGENT
-            if agent_type == "probabilistic_agent":
+            if agent_type == "ProbabilisticAgent":
                # TODO: implement non-random agents and fix this parsing
                settings = agent_cfg.get("agent_settings")
                new_agent = ProbabilisticAgent(
--- a/src/primaite/simulator/network/transmission/network_layer.py
+++ b/src/primaite/simulator/network/transmission/network_layer.py
@@ -16,7 +16,7 @@ class IPProtocol(Enum):
    """

    NONE = "none"
-    """Placeholder for a non-port."""
+    """Placeholder for a non-protocol."""
    TCP = "tcp"
    """Transmission Control Protocol."""
    UDP = "udp"
--- a/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
+++ b/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
@@ -38,9 +38,7 @@ class DataManipulationAttackStage(IntEnum):
 class DataManipulationBot(Application):
    """A bot that simulates a script which performs a SQL injection attack."""

-    # server_ip_address: Optional[IPv4Address] = None
    payload: Optional[str] = None
-    # server_password: Optional[str] = None
    port_scan_p_of_success: float = 0.1
    data_manipulation_p_of_success: float = 0.1

--- a/tests/assets/configs/bad_primaite_session.yaml
+++ b/tests/assets/configs/bad_primaite_session.yaml
@@ -21,7 +21,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/basic_firewall.yaml
+++ b/tests/assets/configs/basic_firewall.yaml
@@ -40,7 +40,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/basic_switched_network.yaml
+++ b/tests/assets/configs/basic_switched_network.yaml
@@ -40,7 +40,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/dmz_network.yaml
+++ b/tests/assets/configs/dmz_network.yaml
@@ -65,7 +65,7 @@ game:
 agents:
  - ref: client_1_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/eval_only_primaite_session.yaml
+++ b/tests/assets/configs/eval_only_primaite_session.yaml
@@ -25,7 +25,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/multi_agent_session.yaml
+++ b/tests/assets/configs/multi_agent_session.yaml
@@ -31,7 +31,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/test_primaite_session.yaml
+++ b/tests/assets/configs/test_primaite_session.yaml
@@ -29,7 +29,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/assets/configs/train_only_primaite_session.yaml
+++ b/tests/assets/configs/train_only_primaite_session.yaml
@@ -25,7 +25,7 @@ game:
 agents:
  - ref: client_2_green_user
    team: GREEN
-    type: probabilistic_agent
+    type: ProbabilisticAgent
    observation_space:
      type: UC2GreenObservation
    action_space:
--- a/tests/integration_tests/game_layer/test_rewards.py
+++ b/tests/integration_tests/game_layer/test_rewards.py
@@ -41,6 +41,7 @@ def test_WebpageUnavailablePenalty(game_and_agent):


 def test_uc2_rewards(game_and_agent):
+    """Test that the reward component correctly applies a penalty when the selected client cannot reach the database."""
    game, agent = game_and_agent
    agent: ControlledAgent