#2417 Move classes to correct files

src/primaite/game/agent/observations/router_observation.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from gymnasium import spaces
+from gymnasium.core import ObsType
+
+from primaite import getLogger
+from primaite.game.agent.observations.acl_observation import ACLObservation
+from primaite.game.agent.observations.nic_observations import PortObservation
+from primaite.game.agent.observations.observations import AbstractObservation, WhereType
+from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
+
+_LOGGER = getLogger(__name__)
+
+
+class RouterObservation(AbstractObservation, identifier="ROUTER"):
+    """Router observation, provides status information about a router within the simulation environment."""
+
+    class ConfigSchema(AbstractObservation.ConfigSchema):
+        """Configuration schema for RouterObservation."""
+
+        hostname: str
+        """Hostname of the router, used for querying simulation state dictionary."""
+        ports: Optional[List[PortObservation.ConfigSchema]] = None
+        """Configuration of port observations for this router."""
+        num_ports: Optional[int] = None
+        """Number of port observations configured for this router."""
+        acl: Optional[ACLObservation.ConfigSchema] = None
+        """Configuration of ACL observation on this router."""
+        ip_list: Optional[List[str]] = None
+        """List of IP addresses for encoding ACLs."""
+        wildcard_list: Optional[List[str]] = None
+        """List of IP wildcards for encoding ACLs."""
+        port_list: Optional[List[int]] = None
+        """List of ports for encoding ACLs."""
+        protocol_list: Optional[List[str]] = None
+        """List of protocols for encoding ACLs."""
+        num_rules: Optional[int] = None
+        """Number of rules ACL rules to show."""
+
+    def __init__(
+        self,
+        where: WhereType,
+        ports: List[PortObservation],
+        num_ports: int,
+        acl: ACLObservation,
+    ) -> None:
+        """
+        Initialize a router observation instance.
+
+        :param where: Where in the simulation state dictionary to find the relevant information for this router.
+            A typical location for a router might be ['network', 'nodes', <node_hostname>].
+        :type where: WhereType
+        :param ports: List of port observations representing the ports of the router.
+        :type ports: List[PortObservation]
+        :param num_ports: Number of ports for the router.
+        :type num_ports: int
+        :param acl: ACL observation representing the access control list of the router.
+        :type acl: ACLObservation
+        """
+        self.where: WhereType = where
+        self.ports: List[PortObservation] = ports
+        self.acl: ACLObservation = acl
+        self.num_ports: int = num_ports
+
+        while len(self.ports) < num_ports:
+            self.ports.append(PortObservation(where=None))
+        while len(self.ports) > num_ports:
+            self.ports.pop()
+            msg = "Too many ports in router observation. Truncating."
+            _LOGGER.warning(msg)
+
+        self.default_observation = {
+            "PORTS": {i + 1: p.default_observation for i, p in enumerate(self.ports)},
+            "ACL": self.acl.default_observation,
+        }
+
+    def observe(self, state: Dict) -> ObsType:
+        """
+        Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary.
+        :type state: Dict
+        :return: Observation containing the status of ports and ACL configuration of the router.
+        :rtype: ObsType
+        """
+        router_state = access_from_nested_dict(state, self.where)
+        if router_state is NOT_PRESENT_IN_STATE:
+            return self.default_observation
+
+        obs = {}
+        obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
+        obs["ACL"] = self.acl.observe(state)
+        return obs
+
+    @property
+    def space(self) -> spaces.Space:
+        """
+        Gymnasium space object describing the observation space shape.
+
+        :return: Gymnasium space representing the observation space for router status.
+        :rtype: spaces.Space
+        """
+        return spaces.Dict(
+            {"PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}), "ACL": self.acl.space}
+        )
+
+    @classmethod
+    def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> RouterObservation:
+        """
+        Create a router observation from a configuration schema.
+
+        :param config: Configuration schema containing the necessary information for the router observation.
+        :type config: ConfigSchema
+        :param parent_where: Where in the simulation state dictionary to find the information about this router's
+            parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
+        :type parent_where: WhereType, optional
+        :return: Constructed router observation instance.
+        :rtype: RouterObservation
+        """
+        where = parent_where + ["nodes", config.hostname]
+
+        if config.acl is None:
+            config.acl = ACLObservation.ConfigSchema()
+        if config.acl.num_rules is None:
+            config.acl.num_rules = config.num_rules
+        if config.acl.ip_list is None:
+            config.acl.ip_list = config.ip_list
+        if config.acl.wildcard_list is None:
+            config.acl.wildcard_list = config.wildcard_list
+        if config.acl.port_list is None:
+            config.acl.port_list = config.port_list
+        if config.acl.protocol_list is None:
+            config.acl.protocol_list = config.protocol_list
+
+        if config.ports is None:
+            config.ports = [PortObservation.ConfigSchema(port_id=i + 1) for i in range(config.num_ports)]
+
+        ports = [PortObservation.from_config(config=c, parent_where=where) for c in config.ports]
+        acl = ACLObservation.from_config(config=config.acl, parent_where=where)
+        return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl)