PrimAITE/src/primaite/game/session.py

# What do? Be an entry point for using PrimAITE
# 1. parse monoconfig
# 2. craete simulation
# 3. create actors and configure their actions/observations/rewards/ anything else
# 4. Create connection with ARCD GATE
# 5. idk

from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional, Tuple
from gymnasium import spaces
from gymnasium.spaces.utils import flatten, flatten_space, unflatten
from gymnasium.core import ObsType, ActType

import numpy as np

from pydantic import BaseModel

from primaite import getLogger
from primaite.game.agent.GATE_agents import GATERLAgent
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent, RandomAgent
from primaite.game.agent.observations import (
    AclObservation,
    FileObservation,
    FolderObservation,
    ICSObservation,
    LinkObservation,
    NicObservation,
    NodeObservation,
    NullObservation,
    ObservationSpace,
    ServiceObservation,
    UC2BlueObservation,
    UC2GreenObservation,
    UC2RedObservation,
)
from primaite.game.agent.rewards import RewardFunction
from primaite.simulator.network.hardware.base import Link, NIC, Node
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.service import Service

from arcd_gate.client.gate_client import GATEClient, ActType
from numpy import ndarray

_LOGGER = getLogger(__name__)

class PrimaiteGATEClient(GATEClient):
    def  __init__(self, parent_session:"PrimaiteSession", service_port: int = 50000):
        super().__init__(service_port=service_port)
        self.parent_session:"PrimaiteSession" = parent_session

    @property
    def rl_framework(self) -> str:
        return self.parent_session.training_options.rl_framework

    @property
    def rl_algorithm(self) -> str:
        return self.parent_session.training_options.rl_algorithm

    @property
    def seed(self) -> int | None:
        return self.parent_session.training_options.seed

    @property
    def n_learn_episodes(self) -> int:
        return self.parent_session.training_options.n_learn_episodes

    @property
    def n_learn_steps(self) -> int:
        return self.parent_session.training_options.n_learn_steps

    @property
    def n_eval_episodes(self) -> int:
        return self.parent_session.training_options.n_eval_episodes

    @property
    def n_eval_steps(self) -> int:
        return self.parent_session.training_options.n_eval_steps

    @property
    def action_space(self) -> spaces.Space:
        return self.parent_session.rl_agent.action_space.space

    @property
    def observation_space(self) -> spaces.Space:
        print("YEEY0")
        print(flatten_space(spaces.Dict({})))
        print("YEEY1")
        # print(self.parent_session.rl_agent.observation_space.space)
        return flatten_space(self.parent_session.rl_agent.observation_space.space)

    def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, Dict]:
        self.parent_session.rl_agent.most_recent_action = action
        self.parent_session.step()
        obs = self.parent_session.rl_agent.observation_space.observe()
        obs = flatten(self.parent_session.rl_agent.observation_space.space, obs)
        rew = self.parent_session.rl_agent.reward_function.calculate()
        term = False
        trunc = False
        info = {}
        return obs, rew, term, trunc, info


    def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ndarray, Dict]:
        self.parent_session.reset()

    def close(self):
        self.parent_session.close()

class PrimaiteSessionOptions(BaseModel):
    ports: List[str]
    protocols: List[str]

class TrainingOptions(BaseModel):
    rl_framework:str
    rl_algorithm:str
    seed:Optional[int]
    n_learn_episodes:int
    n_learn_steps:int
    n_eval_episodes:int
    n_eval_steps:int



class PrimaiteSession:
    def __init__(self):
        self.simulation: Simulation = Simulation()
        self.agents: List[AbstractAgent] = []
        self.rl_agent: AbstractAgent
        # which of the agents should be used for sending RL data to GATE client?
        self.step_counter: int = 0
        self.episode_counter: int = 0
        self.options: PrimaiteSessionOptions
        self.training_options: TrainingOptions

        self.ref_map_nodes: Dict[str, Node] = {}
        self.ref_map_services: Dict[str, Service] = {}
        self.ref_map_links: Dict[str, Link] = {}
        self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self)

    def start_session(self, opts="TODO..."):
        """Commence the session, this gives the gate client control over the simulation/agent loop."""
        self.gate_client.start()

    def step(self):
        _LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}")
        # currently designed with assumption that all agents act once per step in order

        for agent in self.agents:
            # 3. primaite session asks simulation to provide initial state
            # 4. primate session gives state to all agents
            # 5. primaite session asks agents to produce an action based on most recent state
            _LOGGER.debug(f"Sending simulation state to agent {agent.agent_name}")
            sim_state = self.simulation.describe_state()

            # 6. each agent takes most recent state and converts it to CAOS observation
            agent_obs = agent.convert_state_to_obs(sim_state)

            # 7. meanwhile each agent also takes state and calculates reward
            agent_reward = agent.calculate_reward_from_state(sim_state)

            # 8. each agent takes observation and applies decision rule to observation to create CAOS
            #    action(such as random, rulebased, or send to GATE) (therefore, converting CAOS action
            #    to discrete(40) is only necessary for purposes of RL learning, therefore that bit of
            #    code should live inside of the GATE agent subclass)
            # gets action in CAOS format
            _LOGGER.debug(f"Getting agent action")
            agent_action, action_options = agent.get_action(agent_obs, agent_reward)
            # 9. CAOS action is converted into request (extra information might be needed to enrich
            # the request, this is what the execution definition is there for)
            _LOGGER.debug(f"Formatting agent action {agent_action}")  # maybe too many debug log statements
            agent_request = agent.format_request(agent_action, action_options)

            # 10. primaite session receives the action from the agents and asks the simulation to apply each
            _LOGGER.debug(f"Sending request to simulation: {agent_request}")
            self.simulation.apply_request(agent_request)

        _LOGGER.debug(f"Initiating simulation step {self.step_counter}")
        self.simulation.apply_timestep(self.step_counter)
        self.step_counter += 1

    def reset(self):
        pass

    def close(self):
        pass

    @classmethod
    def from_config(cls, cfg: dict) -> "PrimaiteSession":
        sess = cls()
        sess.options = PrimaiteSessionOptions(
            ports=cfg["game_config"]["ports"],
            protocols=cfg["game_config"]["protocols"],
        )
        sess.training_options = TrainingOptions(**cfg['training_config'])
        sim = sess.simulation
        net = sim.network

        sess.ref_map_nodes: Dict[str, Node] = {}
        sess.ref_map_services: Dict[str, Service] = {}
        sess.ref_map_links: Dict[str, Link] = {}

        nodes_cfg = cfg["simulation"]["network"]["nodes"]
        links_cfg = cfg["simulation"]["network"]["links"]
        for node_cfg in nodes_cfg:
            node_ref = node_cfg["ref"]
            n_type = node_cfg["type"]
            if n_type == "computer":
                new_node = Computer(
                    hostname=node_cfg["hostname"],
                    ip_address=node_cfg["ip_address"],
                    subnet_mask=node_cfg["subnet_mask"],
                    default_gateway=node_cfg["default_gateway"],
                    dns_server=node_cfg["dns_server"],
                )
            elif n_type == "server":
                new_node = Server(
                    hostname=node_cfg["hostname"],
                    ip_address=node_cfg["ip_address"],
                    subnet_mask=node_cfg["subnet_mask"],
                    default_gateway=node_cfg["default_gateway"],
                    dns_server=node_cfg.get("dns_server"),
                )
            elif n_type == "switch":
                new_node = Switch(hostname=node_cfg["hostname"], num_ports=node_cfg.get("num_ports"))
            elif n_type == "router":
                new_node = Router(hostname=node_cfg["hostname"], num_ports=node_cfg.get("num_ports"))
                if "ports" in node_cfg:
                    for port_num, port_cfg in node_cfg["ports"].items():
                        new_node.configure_port(
                            port=port_num, ip_address=port_cfg["ip_address"], subnet_mask=port_cfg["subnet_mask"]
                        )
                if "acl" in node_cfg:
                    for r_num, r_cfg in node_cfg["acl"].items():
                        # excuse the uncommon walrus operator ` := `. It's just here as a shorthand, to avoid repeating
                        # this: 'r_cfg.get('src_port')'
                        # Port/IPProtocol. TODO Refactor
                        new_node.acl.add_rule(
                            action=ACLAction[r_cfg["action"]],
                            src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
                            dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
                            protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
                            src_ip_address=r_cfg.get("ip_address"),
                            dst_ip_address=r_cfg.get("ip_address"),
                            position=r_num,
                        )
            else:
                print("invalid node type")
            if "services" in node_cfg:
                for service_cfg in node_cfg["services"]:
                    service_ref = service_cfg["ref"]
                    service_type = service_cfg["type"]
                    service_types_mapping = {
                        "DNSClient": DNSClient,  # key is equal to the 'name' attr of the service class itself.
                        "DNSServer": DNSServer,
                        "DatabaseClient": DatabaseClient,
                        "DatabaseService": DatabaseService,
                        # 'database_backup': ,
                        "DataManipulationBot": DataManipulationBot,
                        # 'web_browser'
                    }
                    if service_type in service_types_mapping:
                        new_node.software_manager.install(service_types_mapping[service_type])
                        new_service = new_node.software_manager.software[service_type]
                        sess.ref_map_services[service_ref] = new_service
                    else:
                        print(f"service type not found {service_type}")
                    # service-dependent options
                    if service_type == "DatabaseClient":
                        if "options" in service_cfg:
                            opt = service_cfg["options"]
                            if "db_server_ip" in opt:
                                new_service.configure(server_ip_address=IPv4Address(opt["db_server_ip"]))
                    if service_type == "DNSServer":
                        if "options" in service_cfg:
                            opt = service_cfg["options"]
                            if "domain_mapping" in opt:
                                for domain, ip in opt["domain_mapping"].items():
                                    new_service.dns_register(domain, ip)
            if "nics" in node_cfg:
                for nic_num, nic_cfg in node_cfg["nics"].items():
                    new_node.connect_nic(NIC(ip_address=nic_cfg["ip_address"], subnet_mask=nic_cfg["subnet_mask"]))

            net.add_node(new_node)
            new_node.power_on()
            sess.ref_map_nodes[node_ref] = new_node.uuid

        # 2. create links between nodes
        for link_cfg in links_cfg:
            node_a = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_a_ref"]]]
            node_b = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_b_ref"]]]
            if isinstance(node_a, Switch):
                endpoint_a = node_a.switch_ports[link_cfg["endpoint_a_port"]]
            else:
                endpoint_a = node_a.ethernet_port[link_cfg["endpoint_a_port"]]
            if isinstance(node_b, Switch):
                endpoint_b = node_b.switch_ports[link_cfg["endpoint_b_port"]]
            else:
                endpoint_b = node_b.ethernet_port[link_cfg["endpoint_b_port"]]
            new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)
            sess.ref_map_links[link_cfg["ref"]] = new_link.uuid

        # 3. create agents
        game_cfg = cfg["game_config"]
        ports_cfg = game_cfg["ports"]
        protocols_cfg = game_cfg["protocols"]
        agents_cfg = game_cfg["agents"]

        for agent_cfg in agents_cfg:
            agent_ref = agent_cfg["ref"]
            agent_type = agent_cfg["type"]
            action_space_cfg = agent_cfg["action_space"]
            observation_space_cfg = agent_cfg["observation_space"]
            reward_function_cfg = agent_cfg["reward_function"]

            # CREATE OBSERVATION SPACE
            obs_space = ObservationSpace.from_config(observation_space_cfg, sess)

            # CREATE ACTION SPACE
            action_space_cfg["options"]["node_uuids"] = []
            # if a list of nodes is defined, convert them from node references to node UUIDs
            for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}):
                if "node_ref" in action_node_option:
                    node_uuid = sess.ref_map_nodes[action_node_option["node_ref"]]
                    action_space_cfg["options"]["node_uuids"].append(node_uuid)
            # Each action space can potentially have a different list of nodes that it can apply to. Therefore,
            # we will pass node_uuids as a part of the action space config.
            # However, it's not possible to specify the node uuids directly in the config, as they are generated
            # dynamically, so we have to translate node references to uuids before passing this config on.

            if "action_list" in action_space_cfg:
                for action_config in action_space_cfg["action_list"]:
                    if "options" in action_config:
                        if "target_router_ref" in action_config["options"]:
                            _target = action_config["options"]["target_router_ref"]
                            action_config["options"]["target_router_uuid"] = sess.ref_map_nodes[_target]

            action_space = ActionManager.from_config(sess, action_space_cfg)

            # CREATE REWARD FUNCTION
            rew_function = RewardFunction.from_config(reward_function_cfg)

            # CREATE AGENT
            if agent_type == "GreenWebBrowsingAgent":
                # TODO: implement non-random agents and fix this parsing
                new_agent = RandomAgent(
                    agent_name=agent_cfg["ref"],
                    action_space=action_space,
                    observation_space=obs_space,
                    reward_function=rew_function,
                )
                sess.agents.append(new_agent)
            elif agent_type == "GATERLAgent":
                new_agent = RandomAgent(
                    agent_name=agent_cfg["ref"],
                    action_space=action_space,
                    observation_space=obs_space,
                    reward_function=rew_function,
                )
                sess.agents.append(new_agent)
                sess.rl_agent = new_agent
            elif agent_type == "RedDatabaseCorruptingAgent":
                new_agent = RandomAgent(
                    agent_name=agent_cfg["ref"],
                    action_space=action_space,
                    observation_space=obs_space,
                    reward_function=rew_function,
                )
                sess.agents.append(new_agent)
            else:
                print("agent type not found")

        return sess