From f1346ae278e8659ce0ef4cc15daced2ad88d2ff2 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 26 Sep 2023 12:54:56 +0100 Subject: [PATCH] put in agent parsing skeleton --- example_config.yaml | 6 +-- sandbox.ipynb | 48 +++++++++++++------ src/primaite/game/actor/interface.py | 35 -------------- src/primaite/game/agent/GATE_agents.py | 5 ++ .../game/{actor => agent}/__init__.py | 0 src/primaite/game/{actor => agent}/actions.py | 0 src/primaite/game/agent/interface.py | 35 ++++++++++++++ .../game/{actor => agent}/observations.py | 0 src/primaite/game/{actor => agent}/rewards.py | 0 src/primaite/game/agent/scripted_agents.py | 9 ++++ .../game_layer/test_observations.py | 2 +- 11 files changed, 86 insertions(+), 54 deletions(-) delete mode 100644 src/primaite/game/actor/interface.py create mode 100644 src/primaite/game/agent/GATE_agents.py rename src/primaite/game/{actor => agent}/__init__.py (100%) rename src/primaite/game/{actor => agent}/actions.py (100%) create mode 100644 src/primaite/game/agent/interface.py rename src/primaite/game/{actor => agent}/observations.py (100%) rename src/primaite/game/{actor => agent}/rewards.py (100%) create mode 100644 src/primaite/game/agent/scripted_agents.py diff --git a/example_config.yaml b/example_config.yaml index e3871b4a..79cfccac 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -16,7 +16,7 @@ game_config: agents: - ref: client_1_green_user team: GREEN - team: SCRIPTED_GREEN_ + type: GreenWebBrowsingAgent observation_space: null action_space: actions: @@ -40,7 +40,7 @@ game_config: - ref: client_1_data_manipulation_red_bot team: RED - type: SCRIPTED_RED_ + type: RedDatabaseCorruptingAgent observation_space: network: nodes: @@ -220,7 +220,7 @@ simulation: default_gateway: 192.168.1.1 dns_server: 192.168.1.10 nics: - 2: + 2: # unfortunately this number is currently meaningless, they're just added in order and take up the next available slot ip_address: 192.168.10.110 subnet_mask: 255.255.255.0 diff --git a/sandbox.ipynb b/sandbox.ipynb index aa39c3e9..05efcfa2 100644 --- a/sandbox.ipynb +++ b/sandbox.ipynb @@ -39,23 +39,23 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 18, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "2023-09-26 11:47:11,032: Added node bc149bf5-ccc4-4dcd-b419-629ec44b2c9a to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,035: Added node 9cacbaee-33cc-4423-a6c8-fe3dd75b1f87 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,042: Added node d4444d66-7cc3-4cd4-acbd-202cb9fe37ff to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,045: Added node af170371-e99b-42b7-9525-65ca64522539 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,049: Added node d6218f34-a104-469d-a08b-97329ad84c19 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,052: Added node 831a3803-ae65-4cee-a17e-9c1220035bc9 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,055: Added node 1b935654-065d-4cb9-82d9-d67fe3d3304e to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,059: Added node dd181916-076b-4d8a-ab97-a32052624b09 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,064: Added node 3137ab20-1a3c-49f2-8ee5-c862216b2435 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n", - "2023-09-26 11:47:11,067: Added node 6ff8b634-7750-4c6d-8109-abf52514dae5 to Network 2c22989f-8f91-4c61-8be9-1afd733b3e1c\n" + "2023-09-26 12:19:50,895: Added node 0fb262e1-a714-420a-aec7-be37f0deeb75 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,898: Added node 310ca8d7-01e0-401e-b604-705c290e5376 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,900: Added node b3b08f1f-7805-47b2-bdb6-3d83098cd740 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,903: Added node adb37f3e-2307-4123-bff3-01f125883be8 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,906: Added node 1a490716-2ccd-452d-b87e-324d29120b59 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,911: Added node 033460d8-0249-4bdd-aaf0-751b24cc0a1e to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,914: Added node 1e7e4e49-78bf-4031-8372-ee71902720f3 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,916: Added node c9f24a13-e5c8-437b-9234-b0c3f8120e2c to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,920: Added node c881f3ee-2176-493b-a6c2-cad829bf0b6d to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n", + "2023-09-26 12:19:50,922: Added node a3ea75d8-bc2c-4713-92a4-2588b4f43ed6 to Network 9318bac2-d9f4-4e71-bb4c-09ffc573ed1c\n" ] }, { @@ -72,9 +72,12 @@ "\n", "\n", "from typing import Dict\n", + "from primaite.game.agent.interface import AbstractAgent\n", "from primaite.simulator.network.hardware.base import NIC, Link, Node\n", "from primaite.simulator.system.services.service import Service\n", "\n", + "from primaite.game.agent.scripted_agents import GreenWebBrowsingAgent, RedDatabaseCorruptingAgent\n", + "from primaite.game.agent.GATE_agents import GATERLAgent\n", "\n", "class PrimaiteSession:\n", "\n", @@ -90,7 +93,7 @@ " # ref_map_agents: Dict[str, AgentInterface] = {}\n", "\n", "\n", - " game = cls()\n", + " session = cls()\n", " with open(cfg_path, 'r') as file:\n", " conf = yaml.safe_load(file)\n", " \n", @@ -177,6 +180,7 @@ " new_node.connect_nic(NIC(ip_address=nic_cfg['ip_address'], subnet_mask=nic_cfg['subnet_mask']))\n", "\n", " net.add_node(new_node)\n", + " new_node.power_on()\n", " ref_map_nodes[node_ref] = new_node.uuid\n", "\n", " #2. create links between nodes\n", @@ -194,11 +198,25 @@ " new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)\n", " ref_map_links[link_cfg['ref']] = new_link.uuid\n", "\n", - " #2. start/setup simulation objects\n", + " session.simulation = sim\n", " #3. create agents\n", + " game_cfg = conf['game_config']\n", + " ports_cfg = game_cfg['ports']\n", + " protocols_cfg = game_cfg['protocols']\n", + " agents_cfg = game_cfg['agents']\n", + "\n", + " for agent_cfg in agents_cfg:\n", + " agent_ref = agent_cfg['ref']\n", + " agent_type = agent_cfg['type']\n", + " action_space_cfg = agent_cfg['action_space']\n", + " observation_space_cfg = agent_cfg['observation_space']\n", + " reward_function_cfg = agent_cfg['reward_function']\n", + " if agent_type == 'GreenWebBrowsingAgent':\n", + " new_agent = GreenWebBrowsingAgent()\n", + "\n", + "\n", " #4. set up agents' actions and observation spaces.\n", - " game.simulation = sim\n", - " return game\n", + " return session\n", "\n", "s = PrimaiteSession.from_config('example_config.yaml')\n", "# print(s.simulation.describe_state())" diff --git a/src/primaite/game/actor/interface.py b/src/primaite/game/actor/interface.py deleted file mode 100644 index d1245e71..00000000 --- a/src/primaite/game/actor/interface.py +++ /dev/null @@ -1,35 +0,0 @@ -# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' -# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word -# If you disagree, make a comment in the PR review and we can discuss -from abc import ABC, abstractmethod -from typing import Any, Dict, List - -from pydantic import BaseModel - -from primaite.game.actor.actions import ActionSpace -from primaite.game.actor.observations import ObservationSpace -from primaite.game.actor.rewards import RewardFunction - - -class AbstractActor(ABC): - """Base class for scripted and RL agents.""" - - def __init__(self) -> None: - self.action_space = ActionSpace - self.observation_space = ObservationSpace - self.reward_function = RewardFunction - - -class AbstractScriptedActor(AbstractActor): - """Base class for actors which generate their own behaviour.""" - - ... - - -class AbstractPuppetActor(AbstractActor): - """Base class for actors controlled via external messages, such as RL policies.""" - - ... - - -# class AbstractRLActor(AbstractPuppetActor): ?? diff --git a/src/primaite/game/agent/GATE_agents.py b/src/primaite/game/agent/GATE_agents.py new file mode 100644 index 00000000..5bdfebe4 --- /dev/null +++ b/src/primaite/game/agent/GATE_agents.py @@ -0,0 +1,5 @@ +from primaite.game.agent.interface import AbstractGATEAgent + + +class GATERLAgent(AbstractGATEAgent): + ... diff --git a/src/primaite/game/actor/__init__.py b/src/primaite/game/agent/__init__.py similarity index 100% rename from src/primaite/game/actor/__init__.py rename to src/primaite/game/agent/__init__.py diff --git a/src/primaite/game/actor/actions.py b/src/primaite/game/agent/actions.py similarity index 100% rename from src/primaite/game/actor/actions.py rename to src/primaite/game/agent/actions.py diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py new file mode 100644 index 00000000..b1ade94b --- /dev/null +++ b/src/primaite/game/agent/interface.py @@ -0,0 +1,35 @@ +# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent' +# That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word +# If you disagree, make a comment in the PR review and we can discuss +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +from primaite.game.agent.actions import ActionSpace +from primaite.game.agent.observations import ObservationSpace +from primaite.game.agent.rewards import RewardFunction + + +class AbstractAgent(ABC): + """Base class for scripted and RL agents.""" + + def __init__( + self, + action_space: Optional[ActionSpace], + observation_space: Optional[ObservationSpace], + reward_function: Optional[RewardFunction], + ) -> None: + self.action_space: Optional[ActionSpace] = action_space + self.observation_space: Optional[ObservationSpace] = observation_space + self.reward_function: Optional[RewardFunction] = reward_function + + +class AbstractScriptedAgent(AbstractAgent): + """Base class for actors which generate their own behaviour.""" + + ... + + +class AbstractGATEAgent(AbstractAgent): + """Base class for actors controlled via external messages, such as RL policies.""" + + ... diff --git a/src/primaite/game/actor/observations.py b/src/primaite/game/agent/observations.py similarity index 100% rename from src/primaite/game/actor/observations.py rename to src/primaite/game/agent/observations.py diff --git a/src/primaite/game/actor/rewards.py b/src/primaite/game/agent/rewards.py similarity index 100% rename from src/primaite/game/actor/rewards.py rename to src/primaite/game/agent/rewards.py diff --git a/src/primaite/game/agent/scripted_agents.py b/src/primaite/game/agent/scripted_agents.py new file mode 100644 index 00000000..d3becd57 --- /dev/null +++ b/src/primaite/game/agent/scripted_agents.py @@ -0,0 +1,9 @@ +from primaite.game.agent.interface import AbstractScriptedAgent + + +class GreenWebBrowsingAgent(AbstractScriptedAgent): + ... + + +class RedDatabaseCorruptingAgent(AbstractScriptedAgent): + ... diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index 4f778f78..7f20a938 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -1,6 +1,6 @@ from gym import spaces -from primaite.game.actor.observations import FileObservation +from primaite.game.agent.observations import FileObservation from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.sim_container import Simulation