Simplify MalSimulator and create wrappers for other interfaces

.github/workflows/test-pytest.yml

@@ -17,8 +17,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install .
+          pip install ".[ml]"
           pip install pytest
       - name: Test with pytest
         run: |
-          pytest tests
+          pytest tests

.gitignore

@@ -8,6 +8,7 @@
 *.swp
 *.swo
 tmp/
+logs/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -35,12 +35,6 @@ they are a setup for running a simulation. This is how the format looks like:
 lang_file: <path to .mar-archive>
 model_file: <path to json/yml model>
 
-attacker_agent_class: 'BreadthFirstAttacker' | 'DepthFirstAttacker' | 'KeyboardAgent'
-
-# For defender_agent_class, null and False are treated the same - no defender will be used in the simulation
-defender_agent_class: 'BreadthFirstAttacker' | 'DepthFirstAttacker' | 'KeyboardAgent' | null | False
-
-
 # Optionally add rewards for each attack step
 rewards:
   <full name of attack step>: <reward>
@@ -50,17 +44,19 @@ rewards:
   # Data A:read: 100
   ...
 
+# Add entry points to AttackGraph with attacker names
+# and attack step full_names
+agents:
+  'Attacker1':
+    type: 'attacker'
+    agent_class: BreadthFirstAttacker | DepthFirstAttacker | KeyboardAgent | null
+    entry_points:
+    - 'Credentials:6:attemptCredentialsReuse'
 
-# Optionally add entry points to AttackGraph with attacker name and attack step full_names.
-# NOTE: If attacker entry points defined in both model and scenario,
-#       the scenario overrides the ones in the model.
-attacker_entry_points:
-  <attacker name>:
-    - <attack step full name>
+  'Defender1':
+    type: 'defender'
+    agent_class: BreadthFirstDefender | DepthFirstDefender | KeyboardAgent | null
 
-  # example:
-  # 'Attacker1':
-  #   - 'Credentials:6:attemptCredentialsReuse'
 
 # Optionally add observability rules that are applied to AttackGrapNodes
 # to make only certain steps observable

malsim/__init__.py

@@ -17,7 +17,7 @@
 
 import logging
 
-from malsim.wrappers.gym_wrapper import AttackerEnv, DefenderEnv, register_envs
+from malsim.mal_simulator import MalSimulator
 
 """
 MAL Simulator
@@ -29,7 +29,7 @@
 __license__ = "Apache 2.0"
 __docformat__ = "restructuredtext en"
 
-__all__ = ("AttackerEnv", "DefenderEnv", "register_envs")
+__all__ = ["MalSimulator"]
 
 
 # TODO: Make sure logging dir exists and make it configurable (or use same as maltoolbox)

malsim/__main__.py

@@ -0,0 +1,92 @@
+"""CLI to run simulations in MAL Simulator using scenario files"""
+
+from __future__ import annotations
+import argparse
+import logging
+
+from .mal_simulator import MalSimulator
+from .agents import DecisionAgent
+from .scenario import create_simulator_from_scenario
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logging.getLogger().setLevel(logging.INFO)
+
+def run_simulation(sim: MalSimulator, agents: list[dict]):
+    """Run a simulation with agents"""
+
+    sim.reset()
+    total_rewards = {agent_dict['name']: 0 for agent_dict in agents}
+    all_agents_term_or_trunc = False
+
+    logger.info("Starting CLI env simulator.")
+
+    i = 1
+    while not all_agents_term_or_trunc:
+        logger.info("Iteration %s", i)
+        all_agents_term_or_trunc = True
+        actions = {}
+
+        # Select actions for each agent
+        for agent_dict in agents:
+            decision_agent: DecisionAgent = agent_dict.get('agent')
+            agent_name = agent_dict['name']
+            if decision_agent is None:
+                logger.warning(
+                    'Agent "%s" has no decision agent class '
+                    'specified in scenario. Waiting.', agent_name,
+                )
+                continue
+
+            sim_agent_state = sim.agent_states[agent_name]
+            agent_action = decision_agent.get_next_action(sim_agent_state)
+            if agent_action:
+                actions[agent_name] = [agent_action]
+                logger.info(
+                    'Agent "%s" chose action: %s',
+                    agent_name, agent_action.full_name
+                )
+
+        # Perform next step of simulation
+        sim.step(actions)
+
+        for agent_dict in agents:
+            agent_name = agent_dict['name']
+            agent_state = sim.agent_states[agent_name]
+            total_rewards[agent_name] += agent_state.reward
+            if not agent_state.terminated and not agent_state.truncated:
+                all_agents_term_or_trunc = False
+        print("---\n")
+        i += 1
+
+    logger.info("Game Over.")
+
+    # Print total rewards
+    for agent_dict in agents:
+        agent_name = agent_dict['name']
+        print(f'Total reward "{agent_name}"', total_rewards[agent_name])
+
+def main():
+    """Entrypoint function of the MAL Toolbox CLI"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        'scenario_file',
+        type=str,
+        help="Can be found in https://github.com/mal-lang/malsim-scenarios/"
+    )
+    parser.add_argument(
+        '-o', '--output-attack-graph', type=str,
+        help="If set to a path, attack graph will be dumped there",
+    )
+    args = parser.parse_args()
+
+    sim, agents = create_simulator_from_scenario(args.scenario_file)
+
+    if args.output_attack_graph:
+        sim.attack_graph.save_to_file(args.output_attack_graph)
+
+    run_simulation(sim, agents)
+
+
+if __name__ == '__main__':
+    main()

malsim/agents/__init__.py

@@ -0,0 +1,12 @@
+from .decision_agent import DecisionAgent
+from .passive_agent import PassiveAgent
+from .keyboard_input import KeyboardAgent
+from .searchers import BreadthFirstAttacker, DepthFirstAttacker
+
+__all__ = [
+    'PassiveAgent',
+    'DecisionAgent',
+    'KeyboardAgent',
+    'BreadthFirstAttacker',
+    'DepthFirstAttacker'
+]

malsim/agents/decision_agent.py

@@ -0,0 +1,28 @@
+"""A decision agent is a heuristic agent"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
+from abc import ABC, abstractmethod
+
+if TYPE_CHECKING:
+    from ..mal_simulator import MalSimAgentStateView
+    from maltoolbox.attackgraph import AttackGraphNode
+
+class DecisionAgent(ABC):
+
+    @abstractmethod
+    def get_next_action(
+        self,
+        agent_state: MalSimAgentStateView,
+        **kwargs
+    ) -> Optional[AttackGraphNode]:
+        """
+        Select next action the agent will work with.
+
+        Attributes:
+            agent: Current state of and other info about the agent from the simulator
+
+        Returns:
+            The selected action or None if there are no actions to select from.
+        """
+        ...

malsim/agents/keyboard_input.py

@@ -1,20 +1,29 @@
-import numpy as np
+from __future__ import annotations
 import logging
+from typing import TYPE_CHECKING, Optional
 
-AGENT_ATTACKER = "attacker"
-AGENT_DEFENDER = "defender"
+from .decision_agent import DecisionAgent
+from ..mal_simulator import MalSimAgentStateView
+
+if TYPE_CHECKING:
+    from maltoolbox.attackgraph import AttackGraphNode
 
 logger = logging.getLogger(__name__)
 
-null_action = (0, None)
+class KeyboardAgent(DecisionAgent):
+    """An agent that makes decisions by asking user for keyboard input"""
 
+    def __init__(self, _, **kwargs):
+        super().__init__(**kwargs)
+        logger.info("Creating KeyboardAgent")
 
-class KeyboardAgent:
-    def __init__(self, vocab):
-        logger.debug("Create Keyboard agent.")
-        self.vocab = vocab
+    def get_next_action(
+            self,
+            agent_state: MalSimAgentStateView,
+            **kwargs
+        ) -> Optional[AttackGraphNode]:
+        """Compute action from action_surface"""
 
-    def compute_action_from_dict(self, obs: dict, mask: tuple) -> tuple:
         def valid_action(user_input: str) -> bool:
             if user_input == "":
                 return True
@@ -24,40 +33,35 @@ def valid_action(user_input: str) -> bool:
             except ValueError:
                 return False
 
-            try:
-                a = associated_action[action_strings[node]]
-            except IndexError:
-                return False
-
-            if a == 0:
-                return True  # wait is always valid
-            return node < len(available_actions) and node >= 0
+            return 0 <= node <= len(agent_state.action_surface)
 
         def get_action_object(user_input: str) -> tuple:
             node = int(user_input) if user_input != "" else None
-            action = associated_action[action_strings[node]] if user_input != "" else 0
-            return node, action
-
-        available_actions = np.flatnonzero(mask[1])
+            return node
 
-        action_strings = [self.vocab[i] for i in available_actions]
-        associated_action = {i: 1 for i in action_strings}
-        action_strings += ["wait"]
-        associated_action["wait"] = 0
+        if not agent_state.action_surface:
+            print("No actions to pick for defender")
+            return []
 
+        index_to_node = dict(enumerate(agent_state.action_surface))
         user_input = "xxx"
         while not valid_action(user_input):
             print("Available actions:")
-            print("\n".join([f"{i}. {a}" for i, a in enumerate(action_strings)]))
+            print(
+                "\n".join(
+                    [f"{i}. {n.full_name}" for i, n in index_to_node.items()]
+                )
+            )
             print("Enter action or leave empty to wait:")
             user_input = input("> ")
 
             if not valid_action(user_input):
                 print("Invalid action.")
 
-        node, a = get_action_object(user_input)
+        index = get_action_object(user_input)
         print(
-            f"Selected action: {action_strings[node] if node is not None else 'wait'}"
+            f"Selected action: {index_to_node[index].full_name}"
+            if index is not None else 'wait'
         )
 
-        return (a, available_actions[node] if a != 0 else -1)
+        return index_to_node[index] if index is not None else None

malsim/agents/passive_agent.py

@@ -0,0 +1,23 @@
+"""A passive agent that always choose to do nothing"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
+
+from .decision_agent import DecisionAgent
+from ..mal_simulator import MalSimAgentStateView
+
+if TYPE_CHECKING:
+    from ..mal_simulator import MalSimAgentStateView
+    from maltoolbox.attackgraph import AttackGraphNode
+
+class PassiveAgent(DecisionAgent):
+    def __init__(self, *args, **kwargs):
+        ...
+
+    def get_next_action(
+        self,
+        agent_state: MalSimAgentStateView,
+        **kwargs
+    ) -> Optional[AttackGraphNode]:
+        # A passive agent never does anything
+        return None