oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2459 back-sync b8 changes into core

Browse Source
This commit is contained in:
Marek Wolan
2024-04-15 11:50:08 +01:00
parent bce32fb376
commit 0828f70b4c
55 changed files with 1383 additions and 441 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
src/primaite/game/agent/actions.py
View File

@@ -487,7 +487,9 @@ class RouterACLAddRuleAction(AbstractAction):
position: int,
permission: int,
source_ip_id: int,
source_wildcard_id: int,
dest_ip_id: int,
dest_wildcard_id: int,
source_port_id: int,
dest_port_id: int,
protocol_id: int,
@@ -519,7 +521,7 @@ class RouterACLAddRuleAction(AbstractAction):
else:
src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
if source_port_id == 0:
return ["do_nothing"] # invalid formulation
elif source_port_id == 1:
@@ -528,13 +530,14 @@ class RouterACLAddRuleAction(AbstractAction):
src_port = self.manager.get_port_by_idx(source_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if source_ip_id == 0:
if dest_ip_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_ip_id == 1:
dst_ip = "ALL"
else:
dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
if dest_port_id == 0:
return ["do_nothing"] # invalid formulation
@@ -553,8 +556,10 @@ class RouterACLAddRuleAction(AbstractAction):
permission_str,
protocol,
str(src_ip),
src_wildcard,
src_port,
str(dst_ip),
dst_wildcard,
dst_port,
position,
]
@@ -624,7 +629,9 @@ class FirewallACLAddRuleAction(AbstractAction):
position: int,
permission: int,
source_ip_id: int,
source_wildcard_id: int,
dest_ip_id: int,
dest_wildcard_id: int,
source_port_id: int,
dest_port_id: int,
protocol_id: int,
@@ -665,7 +672,7 @@ class FirewallACLAddRuleAction(AbstractAction):
src_port = self.manager.get_port_by_idx(source_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if source_ip_id == 0:
if dest_ip_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_ip_id == 1:
dst_ip = "ALL"
@@ -680,6 +687,8 @@ class FirewallACLAddRuleAction(AbstractAction):
else:
dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
return [
"network",
@@ -692,8 +701,10 @@ class FirewallACLAddRuleAction(AbstractAction):
permission_str,
protocol,
str(src_ip),
src_wildcard,
src_port,
str(dst_ip),
dst_wildcard,
dst_port,
position,
]
@@ -871,7 +882,8 @@ class ActionManager:
max_acl_rules: int = 10, # allows calculating shape
protocols: List[str] = ["TCP", "UDP", "ICMP"], # allow mapping index to protocol
ports: List[str] = ["HTTP", "DNS", "ARP", "FTP", "NTP"], # allow mapping index to port
ip_address_list: List[str] = [], # to allow us to map an index to an ip address.
ip_list: List[str] = [], # to allow us to map an index to an ip address.
wildcard_list: List[str] = [], # to allow mapping from wildcard index to
act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions
) -> None:
"""Init method for ActionManager.
@@ -897,8 +909,8 @@ class ActionManager:
:type protocols: List[str]
:param ports: List of ports that are available in the simulation. Used for calculating action shape.
:type ports: List[str]
:param ip_address_list: List of IP addresses that known to this agent. Used for calculating action shape.
:type ip_address_list: Optional[List[str]]
:param ip_list: List of IP addresses that known to this agent. Used for calculating action shape.
:type ip_list: Optional[List[str]]
:param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
:type act_map: Optional[Dict[int, Dict]]
"""
@@ -959,8 +971,10 @@ class ActionManager:
self.protocols: List[str] = protocols
self.ports: List[str] = ports
self.ip_address_list: List[str] = ip_address_list
self.ip_address_list: List[str] = ip_list
self.wildcard_list: List[str] = wildcard_list
if self.wildcard_list == []:
self.wildcard_list = ["NONE"]
# action_args are settings which are applied to the action space as a whole.
global_action_args = {
"num_nodes": len(self.node_names),
@@ -1195,6 +1209,24 @@ class ActionManager:
raise RuntimeError(msg)
return self.ip_address_list[ip_idx]
def get_wildcard_by_idx(self, wildcard_idx: int) -> str:
"""
Get the IP wildcard corresponding to the given index.
:param ip_idx: The index of the IP wildcard to retrieve.
:type ip_idx: int
:return: The wildcard address.
:rtype: str
"""
if wildcard_idx >= len(self.wildcard_list):
msg = (
f"Error: agent attempted to perform an action on ip wildcard {wildcard_idx} but this"
f" is out of range for its action space. Wildcard list: {self.wildcard_list}"
)
_LOGGER.error(msg)
raise RuntimeError(msg)
return self.wildcard_list[wildcard_idx]
def get_port_by_idx(self, port_idx: int) -> str:
"""
Get the port corresponding to the given index.
@@ -1253,37 +1285,14 @@ class ActionManager:
:return: The constructed ActionManager.
:rtype: ActionManager
"""
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
ip_address_order = cfg["options"].pop("ip_address_order", {})
ip_address_list = []
for entry in ip_address_order:
node_name = entry["node_name"]
nic_num = entry["nic_num"]
node_obj = game.simulation.network.get_node_by_hostname(node_name)
ip_address = node_obj.network_interface[nic_num].ip_address
ip_address_list.append(ip_address)
if not ip_address_list:
node_names = [n["node_name"] for n in cfg.get("nodes", {})]
for node_name in node_names:
node_obj = game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
ip_address_list.append(nic_obj.ip_address)
if "ip_list" not in cfg["options"]:
cfg["options"]["ip_list"] = []
obj = cls(
actions=cfg["action_list"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,
ip_address_list=ip_address_list,
act_map=cfg.get("action_map"),
)

5
src/primaite/game/agent/observations/link_observation.py
View File

@@ -43,7 +43,10 @@ class LinkObservation(AbstractObservation, identifier="LINK"):
"""
link_state = access_from_nested_dict(state, self.where)
if link_state is NOT_PRESENT_IN_STATE:
return self.default_observation
self.where[-1] = "<->".join(self.where[-1].split("<->")[::-1]) # try swapping endpoint A and B
link_state = access_from_nested_dict(state, self.where)
if link_state is NOT_PRESENT_IN_STATE:
return self.default_observation
bandwidth = link_state["bandwidth"]
load = link_state["current_load"]

1
src/primaite/game/agent/observations/observation_manager.py
View File

@@ -189,7 +189,6 @@ class ObservationManager:
"""
if config is None:
return cls(NullObservation())
print(config)
obs_type = config["type"]
obs_class = AbstractObservation._registry[obs_type]
observation = obs_class.from_config(config=obs_class.ConfigSchema(**config["options"]))

1
src/primaite/game/agent/rewards.py
View File

@@ -293,6 +293,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
db_state = access_from_nested_dict(state, self.location_in_state)
if db_state is NOT_PRESENT_IN_STATE or "last_connection_successful" not in db_state:
_LOGGER.debug(f"Can't calculate reward for {self.__class__.__name__}")
return 0.0
last_connection_successful = db_state["last_connection_successful"]
if last_connection_successful is False:
return -1.0

64
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -1,8 +1,13 @@
from typing import Dict, Tuple
import random
from typing import Dict, Optional, Tuple
from gymnasium.core import ObsType
from pydantic import BaseModel
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractScriptedAgent
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction
class RandomAgent(AbstractScriptedAgent):
@@ -19,3 +24,60 @@ class RandomAgent(AbstractScriptedAgent):
:rtype: Tuple[str, Dict]
"""
return self.action_manager.get_action(self.action_manager.space.sample())
class PeriodicAgent(AbstractScriptedAgent):
"""Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""
class Settings(BaseModel):
"""Configuration values for when an agent starts performing actions."""
start_step: int = 20
"The timestep at which an agent begins performing it's actions."
start_variance: int = 5
"Deviation around the start step."
frequency: int = 5
"The number of timesteps to wait between performing actions."
variance: int = 0
"The amount the frequency can randomly change to."
max_executions: int = 999999
"Maximum number of times the agent can execute its action."
def __init__(
self,
agent_name: str,
action_space: ActionManager,
observation_space: ObservationManager,
reward_function: RewardFunction,
settings: Optional[Settings] = None,
) -> None:
"""Initialise PeriodicAgent."""
super().__init__(
agent_name=agent_name,
action_space=action_space,
observation_space=observation_space,
reward_function=reward_function,
)
self.settings = settings or PeriodicAgent.Settings()
self._set_next_execution_timestep(timestep=self.settings.start_step, variance=self.settings.start_variance)
self.num_executions = 0
def _set_next_execution_timestep(self, timestep: int, variance: int) -> None:
"""Set the next execution timestep with a configured random variance.
:param timestep: The timestep when the next execute action should be taken.
:type timestep: int
:param variance: Uniform random variance applied to the timestep
:type variance: int
"""
random_increment = random.randint(-variance, variance)
self.next_execution_timestep = timestep + random_increment
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
"""Do nothing, unless the current timestep is the next execution timestep, in which case do the action."""
if timestep == self.next_execution_timestep and self.num_executions < self.settings.max_executions:
self.num_executions += 1
self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance)
return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
return "DONOTHING", {}

78
src/primaite/game/agent/scripted_agents/tap001.py Normal file
View File

@@ -0,0 +1,78 @@
import random
from typing import Dict, Tuple
from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent
class TAP001(AbstractScriptedAgent):
"""
TAP001 | Mobile Malware -- Ransomware Variant.
Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application)
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setup_agent()
next_execution_timestep: int = 0
starting_node_idx: int = 0
installed: bool = False
def _set_next_execution_timestep(self, timestep: int) -> None:
"""Set the next execution timestep with a configured random variance.
:param timestep: The timestep to add variance to.
"""
random_timestep_increment = random.randint(
-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance
)
self.next_execution_timestep = timestep + random_timestep_increment
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
"""Waits until a specific timestep, then attempts to execute the ransomware application.
This application acts a wrapper around the kill-chain, similar to green-analyst and
the previous UC2 data manipulation bot.
:param obs: Current observation for this agent.
:type obs: ObsType
:param timestep: The current simulation timestep, used for scheduling actions
:type timestep: int
:return: Action formatted in CAOS format
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
return "DONOTHING", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
if not self.installed:
self.installed = True
return "NODE_APPLICATION_INSTALL", {
"node_id": self.starting_node_idx,
"application_name": "RansomwareScript",
"ip_address": self.ip_address,
}
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0}
def setup_agent(self) -> None:
"""Set the next execution timestep when the episode resets."""
self._select_start_node()
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
for n, act in self.action_manager.action_map.items():
if not act[0] == "NODE_APPLICATION_INSTALL":
continue
if act[1]["node_id"] == self.starting_node_idx:
self.ip_address = act[1]["ip_address"]
return
raise RuntimeError("TAP001 agent could not find database server ip address in action map")
def _select_start_node(self) -> None:
"""Set the starting starting node of the agent to be a random node from this agent's action manager."""
# we are assuming that every node in the node manager has a data manipulation application at idx 0
num_nodes = len(self.action_manager.node_names)
self.starting_node_idx = random.randint(0, num_nodes - 1)

46
src/primaite/game/game.py
View File

@@ -11,7 +11,10 @@ from primaite.game.agent.observations.observation_manager import ObservationMana
from primaite.game.agent.rewards import RewardFunction, SharedReward
from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
from primaite.game.agent.scripted_agents.tap001 import TAP001
from primaite.game.science import graph_has_cycle, topological_sort
from primaite.simulator.network.airspace import AIR_SPACE
from primaite.simulator.network.hardware.base import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
@@ -26,6 +29,7 @@ from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
@@ -43,6 +47,7 @@ APPLICATION_TYPES_MAPPING = {
"DatabaseClient": DatabaseClient,
"DataManipulationBot": DataManipulationBot,
"DoSBot": DoSBot,
"RansomwareScript": RansomwareScript,
}
"""List of available applications that can be installed on nodes in the PrimAITE Simulation."""
@@ -128,6 +133,8 @@ class PrimaiteGame:
"""
_LOGGER.debug(f"Stepping. Step counter: {self.step_counter}")
self.pre_timestep()
if self.step_counter == 0:
state = self.get_sim_state()
for agent in self.agents.values():
@@ -172,6 +179,10 @@ class PrimaiteGame:
response=response,
)
def pre_timestep(self) -> None:
"""Apply any pre-timestep logic that helps make sure we have the correct observations."""
self.simulation.pre_timestep(self.step_counter)
def advance_timestep(self) -> None:
"""Advance timestep."""
self.step_counter += 1
@@ -211,6 +222,7 @@ class PrimaiteGame:
:return: A PrimaiteGame object.
:rtype: PrimaiteGame
"""
AIR_SPACE.clear()
game = cls()
game.options = PrimaiteGameOptions(**cfg["game"])
game.save_step_metadata = cfg.get("io_settings", {}).get("save_step_metadata") or False
@@ -268,6 +280,9 @@ class PrimaiteGame:
hostname=node_cfg["hostname"],
ip_address=node_cfg["ip_address"],
subnet_mask=node_cfg["subnet_mask"],
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()],
)
else:
msg = f"invalid node type {n_type} in config"
@@ -339,6 +354,19 @@ class PrimaiteGame:
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
data_manipulation_p_of_success=float(opt.get("data_manipulation_p_of_success", "0.1")),
)
elif application_type == "RansomwareScript":
if "options" in application_cfg:
opt = application_cfg["options"]
new_application.configure(
server_ip_address=IPv4Address(opt.get("server_ip")),
server_password=opt.get("server_password"),
payload=opt.get("payload", "ENCRYPT"),
c2_beacon_p_of_success=float(opt.get("c2_beacon_p_of_success", "0.5")),
target_scan_p_of_success=float(opt.get("target_scan_p_of_success", "0.1")),
ransomware_encrypt_p_of_success=float(
opt.get("ransomware_encrypt_p_of_success", "0.1")
),
)
elif application_type == "DatabaseClient":
if "options" in application_cfg:
opt = application_cfg["options"]
@@ -423,6 +451,15 @@ class PrimaiteGame:
reward_function=reward_function,
settings=settings,
)
elif agent_type == "PeriodicAgent":
settings = PeriodicAgent.Settings(**agent_cfg.get("settings", {}))
new_agent = PeriodicAgent(
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=reward_function,
settings=settings,
)
elif agent_type == "ProxyAgent":
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
new_agent = ProxyAgent(
@@ -443,6 +480,15 @@ class PrimaiteGame:
reward_function=reward_function,
agent_settings=agent_settings,
)
elif agent_type == "TAP001":
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
new_agent = TAP001(
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=reward_function,
agent_settings=agent_settings,
)
else:
msg = f"Configuration error: {agent_type} is not a valid agent type."
_LOGGER.error(msg)