Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add heuristic defender decision agent 'ShutdownCompromisedMachinesDef…
Browse files Browse the repository at this point in the history
…ender'
mrkickling committed Jan 22, 2025

Verified

This commit was signed with the committer’s verified signature.
1 parent bf769ab commit c0593c7
Showing 3 changed files with 52 additions and 2 deletions.
4 changes: 3 additions & 1 deletion malsim/agents/__init__.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
from .decision_agent import PassiveAgent, DecisionAgent
from .keyboard_input import KeyboardAgent
from .searchers import BreadthFirstAttacker, DepthFirstAttacker
from .heuristic_agent import ShutdownCompromisedMachinesDefender

__all__ = [
'PassiveAgent',
'DecisionAgent',
'KeyboardAgent',
'BreadthFirstAttacker',
'DepthFirstAttacker'
'DepthFirstAttacker',
'ShutdownCompromisedMachinesDefender'
]
46 changes: 46 additions & 0 deletions malsim/agents/heuristic_agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
from __future__ import annotations
from typing import Optional, TYPE_CHECKING
import logging

import numpy as np

from .decision_agent import DecisionAgent

if TYPE_CHECKING:
from maltoolbox.attackgraph import AttackGraphNode
from ..sims.mal_sim_agent_state import MalSimAgentStateView

logger = logging.getLogger(__name__)

class ShutdownCompromisedMachinesDefender(DecisionAgent):
"""A defender that defends compromised assets using notPresent"""

def __init__(self, agent_config, **_):
# Seed and rng not currently used
seed = (
agent_config["seed"]
if agent_config.get("seed")
else np.random.SeedSequence().entropy
)
self.rng = (
np.random.default_rng(seed)
if agent_config.get("randomize")
else None
)

def get_next_action(
self, agent: MalSimAgentStateView, **kwargs
) -> Optional[AttackGraphNode]:

"""Return an action that disables a compromised asset"""

selected_node = None
for node in agent.action_surface:

# Child of a defense node is compromised -> enable the defense
# TODO: optionally randomize order so not always same.
for child_node in node.children:
if child_node.is_compromised():
return node

return selected_node
4 changes: 3 additions & 1 deletion malsim/scenario.py
Original file line number Diff line number Diff line change
@@ -23,7 +23,8 @@
BreadthFirstAttacker,
DepthFirstAttacker,
KeyboardAgent,
PassiveAgent
PassiveAgent,
ShutdownCompromisedMachinesDefender
)

from .sims.mal_simulator import AgentType, MalSimulator
@@ -33,6 +34,7 @@
'BreadthFirstAttacker': BreadthFirstAttacker,
'KeyboardAgent': KeyboardAgent,
'PassiveAgent': PassiveAgent,
'ShutdownCompromisedMachinesAgent': ShutdownCompromisedMachinesDefender
}

deprecated_fields = [

0 comments on commit c0593c7

Please sign in to comment.