typing

Author: mrkickling

malsim/mal_simulator.py

@@ -230,7 +230,7 @@ def from_scenario(
     ) -> MalSimulator:
         """Create a MalSimulator object from a Scenario"""
 
-        def register_agent_dict(agent_config: dict):
+        def register_agent_dict(agent_config: dict[str, Any]) -> None:
             """Register an agent specified in a dictionary"""
             logger.warning(
                 "Having agent configs in dictionaries will be deprecated in "
@@ -245,7 +245,7 @@ def register_agent_dict(agent_config: dict):
             elif agent_config['type'] == AgentType.DEFENDER:
                 sim.register_defender(agent_config['name'])
 
-        def register_agent_config(agent_config: AgentConfig):
+        def register_agent_config(agent_config: AgentConfig) -> None:
             """Register an agent config in simulator"""
             if isinstance(agent_config, AttackerAgentConfig):
                 sim.register_attacker(

malsim/scenario.py

@@ -10,7 +10,7 @@
         - attacker_class
         - defender_class
 """
-
+from __future__ import annotations
 import os
 from dataclasses import dataclass, asdict
 from typing import Any, Optional, TextIO
@@ -110,7 +110,7 @@ def __init__(
         self,
         lang_file: str,
         agents: dict[str, Any],
-        model_dict: Optional[dict] = None,
+        model_dict: Optional[dict[str, Any]] = None,
         model_file: Optional[str] = None,
         rewards: Optional[dict[str, Any]] = None,
         false_positive_rates: Optional[dict[str, Any]] = None,
@@ -157,7 +157,7 @@ def __init__(
             self.attack_graph, is_actionable or {}
         )
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, Any]:
         assert self._lang_file, "Can not save scenario to file if lang file was not given"
         scenario_dict = {
             # 'version': ?
@@ -179,14 +179,14 @@ def to_dict(self):
 
         return scenario_dict
 
-    def save_to_file(self, file_path):
+    def save_to_file(self, file_path: str) -> None:
         save_scenario_dict(self.to_dict(), file_path)
 
     @classmethod
-    def from_dict(cls, scenario_dict):
+    def from_dict(cls, scenario_dict: dict[str, Any]) -> Scenario:
         return Scenario(
-            lang_file=scenario_dict.get('lang_file'),
-            agents=scenario_dict.get('agents'),
+            lang_file=scenario_dict['lang_file'],
+            agents=scenario_dict['agents'],
             model_dict=scenario_dict.get('model'),
             model_file=scenario_dict.get('model_file'),
             rewards=scenario_dict.get('rewards'),
@@ -197,7 +197,7 @@ def from_dict(cls, scenario_dict):
         )
 
     @classmethod
-    def load_from_file(cls, scenario_file):
+    def load_from_file(cls, scenario_file: str) -> Scenario:
         scenario_dict = load_scenario_dict(scenario_file)
         _validate_scenario_dict(scenario_dict)
         return cls.from_dict(scenario_dict)
@@ -439,7 +439,6 @@ def get_entry_point_nodes(
 def load_simulator_agents(
         attack_graph: AttackGraph,
         scenario_agents: dict[str, Any],
-        as_dicts=False
     ) -> list[AgentConfig]:
     """Load agents to be registered in MALSimulator
 

tests/envs/test_vectorized_obs_mal_simulator.py

@@ -176,8 +176,6 @@ def test_create_blank_observation_observability_given(
         elif node.model_asset and node.model_asset.name == 'User:3' and node.name in ('phishing'):
             assert observable
         else:
-            if observable:
-                breakpoint()
             assert not observable
 
 def test_create_blank_observation_actionability_given(

tests/test_scenario.py

@@ -6,6 +6,7 @@
 
 from maltoolbox.attackgraph import create_attack_graph
 from malsim.scenario import (
+    AttackerAgentConfig,
     apply_scenario_node_property,
     load_scenario,
     _validate_scenario_node_property_config
@@ -53,6 +54,7 @@ def test_load_scenario() -> None:
 
     # Verify attacker entrypoint was added
     attack_step = get_node(scenario.attack_graph, 'OS App:fullAccess')
+    assert isinstance(scenario.agents[0], AttackerAgentConfig)
     assert attack_step in scenario.agents[0].entry_points
 
     assert isinstance(scenario.agents[0].agent, BreadthFirstAttacker)
@@ -232,7 +234,7 @@ def test_apply_scenario_observability() -> None:
 
     # Apply observability rules
     observable = apply_scenario_node_property(
-        scenario.attack_graph, observability_rules, default_value = 0
+        scenario.attack_graph, observability_rules,
     )
 
     # Make sure all attack steps are observable
@@ -261,53 +263,46 @@ def test_apply_scenario_observability_faulty() -> None:
         apply_scenario_node_property(
             scenario.attack_graph,
             {'NotAllowedKey': {'Data': ['read', 'write', 'delete']}},
-            default_value = 0
         )
 
     # Correct asset type and attack step
     apply_scenario_node_property(
         scenario.attack_graph,
         {'by_asset_type': { 'Application': ['read']}},
-        default_value = 0
     )
 
     # Wrong asset type in rule asset type to step dict
     with pytest.raises(AssertionError):
         apply_scenario_node_property(
             scenario.attack_graph,
             {'by_asset_type': {'NonExistingType': ['read']}},
-            default_value = 0
         )
 
     # Wrong attack step name in rule asset type to step dict
     with pytest.raises(AssertionError):
         apply_scenario_node_property(
             scenario.attack_graph,
             {'by_asset_type': {'Data': ['nonExistingAttackStep']}},
-            default_value = 0
         )
 
     # Correct asset name and attack step
     apply_scenario_node_property(
         scenario.attack_graph,
         {'by_asset_name': { 'OS App': ['read']}},
-        default_value = 0
     )
 
     # Wrong asset name in rule asset name to step dict
     with pytest.raises(AssertionError):
         apply_scenario_node_property(
             scenario.attack_graph,
             {'by_asset_name': { 'NonExistingName': ['read']}},
-            default_value = 0
         )
 
     # Wrong attack step name in rule asset name to step dict
     with pytest.raises(AssertionError):
         apply_scenario_node_property(
             scenario.attack_graph,
             {'by_asset_name': {'OS App': ['nonExistingAttackStep']}},
-            default_value = 0
         )