#2825 Updated all tests to use a constants rather than a magic num.

2025-02-27 14:11:27 +00:00
parent 72a06ea001
commit 6a0d28bf5f
7 changed files with 80 additions and 48 deletions
--- a/tests/e2e_integration_tests/threat_actor_profiles/test_kill_chain_methods.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_kill_chain_methods.py
@@ -19,6 +19,7 @@ from primaite.session.environment import PrimaiteGymEnv
 START_STEP = 1  # The starting step of the agent.
 FREQUENCY = 2  # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
 VARIANCE = 0  # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
+ATTACK_AGENT_INDEX = 32


 def uc7_tap003_env() -> PrimaiteGymEnv:
@@ -26,9 +27,9 @@ def uc7_tap003_env() -> PrimaiteGymEnv:
    with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
    env = PrimaiteGymEnv(env_config=cfg)
    return env

@@ -38,9 +39,9 @@ def uc7_tap001_env() -> PrimaiteGymEnv:
    with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
    env = PrimaiteGymEnv(env_config=cfg)
    return env

--- a/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_repeat.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_repeat.py
@@ -19,6 +19,7 @@ from primaite.session.environment import PrimaiteGymEnv
 START_STEP = 1  # The starting step of the agent.
 FREQUENCY = 2  # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
 VARIANCE = 0  # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
+ATTACK_AGENT_INDEX = 32


 def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
@@ -26,13 +27,19 @@ def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
    with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = kwargs["repeat_kill_chain_stages"]
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PROPAGATE"]["probability"] = kwargs["propagate_probability"]
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["probability"] = kwargs["payload_probability"]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = kwargs[
+            "repeat_kill_chain_stages"
+        ]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PROPAGATE"]["probability"] = kwargs[
+            "propagate_probability"
+        ]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PAYLOAD"]["probability"] = kwargs[
+            "payload_probability"
+        ]
    env = PrimaiteGymEnv(env_config=cfg)
    return env

--- a/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_stages.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_stages.py
@@ -26,6 +26,7 @@ REPEAT_KILL_CHAIN = False  # Should the TAP repeat the kill chain after success/
 REPEAT_KILL_CHAIN_STAGES = False  # Should the TAP restart from it's previous stage on failure?
 KILL_CHAIN_PROBABILITY = 1  # Blank probability for agent 'success'
 DATA_EXFIL = True  # Data exfiltration on the payload stage is enabled.
+ATTACK_AGENT_INDEX = 32


 def uc7_tap001_env() -> PrimaiteGymEnv:
@@ -33,14 +34,18 @@ def uc7_tap001_env() -> PrimaiteGymEnv:
    with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["probability"] = KILL_CHAIN_PROBABILITY
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PROPAGATE"]["probability"] = KILL_CHAIN_PROBABILITY
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["exfiltrate"] = DATA_EXFIL
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PAYLOAD"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PROPAGATE"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PAYLOAD"]["exfiltrate"] = DATA_EXFIL
    env = PrimaiteGymEnv(env_config=cfg)
    return env

--- a/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_propagate_stage.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_tap001_propagate_stage.py
@@ -24,7 +24,8 @@ FREQUENCY = 2  # The frequency of kill chain stage progression (E.g it's next at
 VARIANCE = 0  # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
 REPEAT_KILL_CHAIN = False  # Should the TAP repeat the kill chain after success/failure?
 REPEAT_KILL_CHAIN_STAGES = False  # Should the TAP restart from it's previous stage on failure?
-KILL_CHAIN_PROBABILITY = 1  # Blank probability for agent 'success'
+KILL_CHAIN_PROBABILITY = 1  # Blank probability for agent 'success's.
+ATTACK_AGENT_INDEX = 32


 def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
@@ -32,7 +33,7 @@ def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
    with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        agent_cfg = cfg["agents"][32]["agent_settings"]
+        agent_cfg = cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]
        agent_cfg["start_step"] = START_STEP
        agent_cfg["frequency"] = FREQUENCY
        agent_cfg["variance"] = VARIANCE
@@ -135,6 +136,6 @@ def test_tap001_kill_chain_stage_PROPAGATE_repeat_scan():
    while tap001.current_kill_chain_stage == MobileMalwareKillChain.PROPAGATE:
        env.step(0)

-    # As the given network_address does not contain the target, we should failed because the maximum amount of scan attempts has been reached
-    assert tap001.scans_complete == 20
+    # As the given network_address does not contain the target, we should fail because the maximum amount of scan attempts has been reached
+    assert tap001.scans_complete == scan_attempts
    assert tap001.current_kill_chain_stage == MobileMalwareKillChain.FAILED
--- a/tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_repeat.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_repeat.py
@@ -19,23 +19,30 @@ from primaite.session.environment import PrimaiteGymEnv
 START_STEP = 1  # The starting step of the agent.
 FREQUENCY = 2  # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
 VARIANCE = 0  # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
+ATTACK_AGENT_INDEX = 32


 def uc7_tap003_env(**kwargs) -> PrimaiteGymEnv:
-    """Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
+    """Setups the UC7 TAP003 Game with a 1 timestep start_step, frequency of 2 and probabilities set to 1 as well"""
    with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = kwargs["repeat_kill_chain_stages"]
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = kwargs[
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = kwargs[
+            "repeat_kill_chain_stages"
+        ]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = kwargs[
            "manipulation_probability"
        ]
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = kwargs["access_probability"]
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs["planning_probability"]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = kwargs[
+            "access_probability"
+        ]
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs[
+            "planning_probability"
+        ]
    env = PrimaiteGymEnv(env_config=cfg)
    return env

--- a/tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_stages.py
+++ b/tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_stages.py
@@ -23,22 +23,31 @@ VARIANCE = 0  # The timestep variance between kill chain progression (E.g Next t
 REPEAT_KILL_CHAIN = False  # Should the TAP repeat the kill chain after success/failure?
 REPEAT_KILL_CHAIN_STAGES = False  # Should the TAP restart from it's previous stage on failure?
 KILL_CHAIN_PROBABILITY = 1  # Blank probability for agent 'success'
+ATTACK_AGENT_INDEX = 32


 def uc7_tap003_env() -> PrimaiteGymEnv:
-    """Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
+    """Setups the UC7 TAP003 Game with a 1 timestep start_step, frequency of 2 and probabilities set to 1 as well"""
    with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
        cfg["io_settings"]["save_sys_logs"] = False
-        cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
-        cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
-        cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
-        cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = KILL_CHAIN_PROBABILITY
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = KILL_CHAIN_PROBABILITY
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = KILL_CHAIN_PROBABILITY
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["EXPLOIT"]["probability"] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["EXPLOIT"][
+            "probability"
+        ] = KILL_CHAIN_PROBABILITY
    env = PrimaiteGymEnv(env_config=cfg)
    return env

@@ -215,7 +224,7 @@ def test_tap003_kill_chain_stage_exploit():
    assert tap003.current_kill_chain_stage.name == InsiderKillChain.EXPLOIT.name

    # Testing that the stage successfully impacted the simulation - Malicious ACL Added:
-    for _ in range(32):
+    for _ in range(ATTACK_AGENT_INDEX):
        env.step(0)

    # Tests that the ACL has been added and that the action is deny.
--- a/tests/integration_tests/configuration_file_parsing/threat_actor_profile_settings/test_kill_chain_settings.py
+++ b/tests/integration_tests/configuration_file_parsing/threat_actor_profile_settings/test_kill_chain_settings.py
@@ -8,13 +8,15 @@ from primaite.config.load import _EXAMPLE_CFG
 from primaite.game.agent.scripted_agents.TAP003 import TAP003
 from primaite.game.game import PrimaiteGame

+ATTACK_AGENT_INDEX = 32
+

 def test_tap003_kill_chain_settings_load_config():
    with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
        cfg = yaml.safe_load(uc7_config)
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = 0.5
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = 0.5
-        cfg["agents"][32]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = 0.5
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = 0.5
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = 0.5
+        cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = 0.5
    game = PrimaiteGame.from_config(cfg)
    tap: TAP003 = game.agents["attacker"]
    kill_chain = tap.config.agent_settings.kill_chain