PrimAITE/src/primaite/game/agent/observations/node_observations.py

from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, Dict, Iterable, List, Optional

from gymnasium import spaces
from gymnasium.core import ObsType

from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE

_LOGGER = getLogger(__name__)

WhereType = Iterable[str | int] | None


class ServiceObservation(AbstractObservation, identifier="SERVICE"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        service_name: str

    def __init__(self, where: WhereType)->None:
        self.where = where
        self.default_observation = {"operating_status": 0, "health_status": 0}

    def observe(self, state: Dict) -> Any:
        service_state = access_from_nested_dict(state, self.where)
        if service_state is NOT_PRESENT_IN_STATE:
            return self.default_observation
        return {
            "operating_status": service_state["operating_state"],
            "health_status": service_state["health_state_visible"],
        }

    @property
    def space(self) -> spaces.Space:
        """Gymnasium space object describing the observation space shape."""
        return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
        return cls(where=parent_where+["services", config.service_name])


class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        application_name: str

    def __init__(self, where: WhereType)->None:
        self.where = where
        self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}

    def observe(self, state: Dict) -> Any:
        # raise NotImplementedError("TODO NUM EXECUTIONS NEEDS TO BE CONVERTED TO A CATEGORICAL")
        application_state = access_from_nested_dict(state, self.where)
        if application_state is NOT_PRESENT_IN_STATE:
            return self.default_observation
        return {
            "operating_status": application_state["operating_state"],
            "health_status": application_state["health_state_visible"],
            "num_executions": application_state["num_executions"],
        }

    @property
    def space(self) -> spaces.Space:
        """Gymnasium space object describing the observation space shape."""
        return spaces.Dict({
            "operating_status": spaces.Discrete(7),
            "health_status": spaces.Discrete(5),
            "num_executions": spaces.Discrete(4)
            })

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ApplicationObservation:
        return cls(where=parent_where+["applications", config.application_name])


class FileObservation(AbstractObservation, identifier="FILE"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        file_name: str
        include_num_access: Optional[bool] = None

    def __init__(self, where: WhereType, include_num_access: bool)->None:
        self.where: WhereType = where
        self.include_num_access :bool = include_num_access

        self.default_observation: ObsType = {"health_status": 0}
        if self.include_num_access:
            self.default_observation["num_access"] = 0

    def observe(self, state: Dict) -> Any:
        file_state = access_from_nested_dict(state, self.where)
        if file_state is NOT_PRESENT_IN_STATE:
            return self.default_observation
        obs = {"health_status": file_state["visible_status"]}
        if self.include_num_access:
            obs["num_access"] = file_state["num_access"]
            # raise NotImplementedError("TODO: need to fix num_access to use thresholds instead of raw value.")
        return obs

    @property
    def space(self) -> spaces.Space:
        space = {"health_status": spaces.Discrete(6)}
        if self.include_num_access:
            space["num_access"] = spaces.Discrete(4)
        return spaces.Dict(space)

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> FileObservation:
        return cls(where=parent_where+["files", config.file_name], include_num_access=config.include_num_access)


class FolderObservation(AbstractObservation, identifier="FOLDER"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        folder_name: str
        files: List[FileObservation.ConfigSchema] = []
        num_files : Optional[int] = None
        include_num_access : Optional[bool] = None

    def __init__(self, where: WhereType, files: Iterable[FileObservation], num_files: int, include_num_access: bool)->None:
        self.where: WhereType = where

        self.files: List[FileObservation] = files
        while len(self.files) < num_files:
            self.files.append(FileObservation(where=None,include_num_access=include_num_access))
        while len(self.files) > num_files:
            truncated_file = self.files.pop()
            msg = f"Too many files in folder observation. Truncating file {truncated_file}"
            _LOGGER.warning(msg)

        self.default_observation = {
            "health_status": 0,
            "FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)},
        }

    def observe(self, state: Dict) -> Any:
        folder_state = access_from_nested_dict(state, self.where)
        if folder_state is NOT_PRESENT_IN_STATE:
            return self.default_observation

        health_status = folder_state["health_status"]

        obs = {}

        obs["health_status"] = health_status
        obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)}

        return obs

    @property
    def space(self) -> spaces.Space:
        """Gymnasium space object describing the observation space shape.

        :return: Gymnasium space
        :rtype: spaces.Space
        """
        return spaces.Dict(
            {
                "health_status": spaces.Discrete(6),
                "FILES": spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)}),
            }
        )
    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> FolderObservation:
        where = parent_where + ["folders", config.folder_name]

        #pass down shared/common config items
        for file_config in config.files:
            file_config.include_num_access = config.include_num_access

        files = [FileObservation.from_config(config=f, parent_where = where) for f in config.files]
        return cls(where=where, files=files, num_files=config.num_files, include_num_access=config.include_num_access)


class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        nic_num: int
        include_nmne: Optional[bool] = None


    def __init__(self, where: WhereType, include_nmne: bool)->None:
        self.where = where
        self.include_nmne : bool = include_nmne

        self.default_observation: ObsType = {"nic_status": 0}
        if self.include_nmne:
            self.default_observation.update({"NMNE":{"inbound":0, "outbound":0}})

    def observe(self, state: Dict) -> Any:
        # raise NotImplementedError("TODO: CATEGORISATION")
        nic_state = access_from_nested_dict(state, self.where)

        if nic_state is NOT_PRESENT_IN_STATE:
            return self.default_observation

        obs = {"nic_status": 1 if nic_state["enabled"] else 2}
        if self.include_nmne:
            obs.update({"NMNE": {}})
            direction_dict = nic_state["nmne"].get("direction", {})
            inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
            inbound_count = inbound_keywords.get("*", 0)
            outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {})
            outbound_count = outbound_keywords.get("*", 0)
            obs["NMNE"]["inbound"] = self._categorise_mne_count(inbound_count - self.nmne_inbound_last_step)
            obs["NMNE"]["outbound"] = self._categorise_mne_count(outbound_count - self.nmne_outbound_last_step)
            self.nmne_inbound_last_step = inbound_count
            self.nmne_outbound_last_step = outbound_count
        return obs


    @property
    def space(self) -> spaces.Space:
        space = spaces.Dict({"nic_status": spaces.Discrete(3)})

        if self.include_nmne:
            space["NMNE"] = spaces.Dict({"inbound": spaces.Discrete(4), "outbound": spaces.Discrete(4)})

        return space

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> NICObservation:
        return cls(where = parent_where+["NICs", config.nic_num], include_nmne=config.include_nmne)


class HostObservation(AbstractObservation, identifier="HOST"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        hostname: str
        services: List[ServiceObservation.ConfigSchema] = []
        applications: List[ApplicationObservation.ConfigSchema] = []
        folders: List[FolderObservation.ConfigSchema] = []
        network_interfaces: List[NICObservation.ConfigSchema] = []
        num_services: Optional[int] = None
        num_applications: Optional[int] = None
        num_folders: Optional[int] = None
        num_files: Optional[int] = None
        num_nics: Optional[int] = None
        include_nmne: Optional[bool] = None
        include_num_access: Optional[bool] = None

    def __init__(self,
                 where: WhereType,
                 services:List[ServiceObservation],
                 applications:List[ApplicationObservation],
                 folders:List[FolderObservation],
                 network_interfaces:List[NICObservation],
                 num_services: int,
                 num_applications: int,
                 num_folders: int,
                 num_files: int,
                 num_nics: int,
                 include_nmne: bool,
                 include_num_access: bool
                 )->None:

        self.where : WhereType = where

        # ensure service list has length equal to num_services by truncating or padding
        self.services: List[ServiceObservation] = services
        while len(self.services) < num_services:
            self.services.append(ServiceObservation(where=None))
        while len(self.services) > num_services:
            truncated_service = self.services.pop()
            msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
            _LOGGER.warning(msg)

        # ensure application list has length equal to num_applications by truncating or padding
        self.applications: List[ApplicationObservation] = applications
        while len(self.applications) < num_applications:
            self.applications.append(ApplicationObservation(where=None))
        while len(self.applications) > num_applications:
            truncated_application = self.applications.pop()
            msg = f"Too many applications in Node observation space for node. Truncating application {truncated_application.where}"
            _LOGGER.warning(msg)

        # ensure folder list has length equal to num_folders by truncating or padding
        self.folders: List[FolderObservation] = folders
        while len(self.folders) < num_folders:
            self.folders.append(FolderObservation(where = None, files= [], num_files=num_files, include_num_access=include_num_access))
        while len(self.folders) > num_folders:
            truncated_folder = self.folders.pop()
            msg = f"Too many folders in Node observation space for node. Truncating folder {truncated_folder.where}"
            _LOGGER.warning(msg)

        # ensure network_interface list has length equal to num_network_interfaces by truncating or padding
        self.network_interfaces: List[NICObservation] = network_interfaces
        while len(self.network_interfaces) < num_nics:
            self.network_interfaces.append(NICObservation(where = None, include_nmne=include_nmne))
        while len(self.network_interfaces) > num_nics:
            truncated_nic = self.network_interfaces.pop()
            msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_folder.where}"
            _LOGGER.warning(msg)

        self.default_observation: ObsType = {
            "SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
            "APPLICATIONS": {i + 1: a.default_observation for i, a in enumerate(self.applications)},
            "FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)},
            "NICS": {i + 1: n.default_observation for i, n in enumerate(self.network_interfaces)},
            "operating_status": 0,
            "num_file_creations": 0,
            "num_file_deletions": 0,
        }


    def observe(self, state: Dict) -> Any:
        node_state = access_from_nested_dict(state, self.where)
        if node_state is NOT_PRESENT_IN_STATE:
            return self.default_observation

        obs = {}
        obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
        obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)}
        obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
        obs["operating_status"] = node_state["operating_state"]
        obs["NICS"] = {
            i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces)
        }
        obs["num_file_creations"] = node_state["file_system"]["num_file_creations"]
        obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"]
        return obs

    @property
    def space(self) -> spaces.Space:
        shape = {
            "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
            "APPLICATIONS": spaces.Dict({i + 1: app.space for i, app in enumerate(self.applications)}),
            "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
            "operating_status": spaces.Discrete(5),
            "NICS": spaces.Dict(
                {i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}
            ),
            "num_file_creations" : spaces.Discrete(4),
            "num_file_deletions" : spaces.Discrete(4),
        }
        return spaces.Dict(shape)

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = None ) -> HostObservation:
        if parent_where is None:
            where = ["network", "nodes", config.hostname]
        else:
            where = parent_where + ["nodes", config.hostname]

        #pass down shared/common config items
        for folder_config in config.folders:
            folder_config.include_num_access = config.include_num_access
            folder_config.num_files = config.num_files
        for nic_config in config.network_interfaces:
            nic_config.include_nmne = config.include_nmne

        services = [ServiceObservation.from_config(config=c,parent_where=where) for c in config.services]
        applications = [ApplicationObservation.from_config(config=c, parent_where=where) for c in config.applications]
        folders = [FolderObservation.from_config(config=c, parent_where=where) for c in config.folders]
        nics = [NICObservation.from_config(config=c, parent_where=where) for c in config.network_interfaces]

        return cls(
            where = where,
            services = services,
            applications = applications,
            folders = folders,
            network_interfaces = nics,
            num_services = config.num_services,
            num_applications = config.num_applications,
            num_folders = config.num_folders,
            num_files = config.num_files,
            num_nics = config.num_nics,
            include_nmne = config.include_nmne,
            include_num_access = config.include_num_access,
        )


class PortObservation(AbstractObservation, identifier="PORT"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        port_id : int

    def __init__(self, where: WhereType)->None:
        self.where = where
        self.default_observation: ObsType = {"operating_status" : 0}

    def observe(self, state: Dict) -> Any:
        port_state = access_from_nested_dict(state, self.where)
        if port_state is NOT_PRESENT_IN_STATE:
            return self.default_observation
        return {"operating_status": 1 if port_state["enabled"] else 2 }

    @property
    def space(self) -> spaces.Space:
        return spaces.Dict({"operating_status": spaces.Discrete(3)})

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> PortObservation:
        return cls(where = parent_where + ["NICs", config.port_id])

class ACLObservation(AbstractObservation, identifier="ACL"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        ip_list: Optional[List[IPv4Address]] = None
        wildcard_list: Optional[List[str]] = None
        port_list: Optional[List[int]] = None
        protocol_list: Optional[List[str]] = None
        num_rules: Optional[int] = None

    def __init__(self, where: WhereType, num_rules: int, ip_list: List[IPv4Address], wildcard_list: List[str], port_list: List[int],protocol_list: List[str])->None:
        self.where = where
        self.num_rules: int = num_rules
        self.ip_to_id: Dict[str, int] = {i+2:p for i,p in enumerate(ip_list)}
        self.wildcard_to_id: Dict[str, int] = {i+2:p for i,p in enumerate(wildcard_list)}
        self.port_to_id: Dict[int, int] = {i+2:p for i,p in enumerate(port_list)}
        self.protocol_to_id: Dict[str, int] = {i+2:p for i,p in enumerate(protocol_list)}
        self.default_observation: Dict = {
            i
            + 1: {
                "position": i,
                "permission": 0,
                "source_ip_id": 0,
                "source_wildcard_id": 0,
                "source_port_id": 0,
                "dest_ip_id": 0,
                "dest_wildcard_id": 0,
                "dest_port_id": 0,
                "protocol_id": 0,
            }
            for i in range(self.num_rules)
        }

    def observe(self, state: Dict) -> Any:
        acl_state: Dict = access_from_nested_dict(state, self.where)
        if acl_state is NOT_PRESENT_IN_STATE:
            return self.default_observation
        obs = {}
        acl_items = dict(acl_state.items())
        i = 1  # don't show rule 0 for compatibility reasons.
        while i < self.num_rules + 1:
            rule_state = acl_items[i]
            if rule_state is None:
                obs[i] = {
                    "position": i - 1,
                    "permission": 0,
                    "source_ip_id": 0,
                    "source_wildcard_id": 0,
                    "source_port_id": 0,
                    "dest_ip_id": 0,
                    "dest_wildcard_id": 0,
                    "dest_port_id": 0,
                    "protocol_id": 0,
                }
            else:
                src_ip = rule_state["src_ip_address"]
                src_node_id = self.ip_to_id.get(src_ip, 1)
                dst_ip = rule_state["dst_ip_address"]
                dst_node_ip = self.ip_to_id.get(dst_ip, 1)
                src_wildcard = rule_state["source_wildcard_id"]
                src_wildcard_id = self.wildcard_to_id.get(src_wildcard, 1)
                dst_wildcard = rule_state["dest_wildcard_id"]
                dst_wildcard_id = self.wildcard_to_id.get(dst_wildcard, 1)
                src_port = rule_state["source_port_id"]
                src_port_id = self.port_to_id.get(src_port, 1)
                dst_port = rule_state["dest_port_id"]
                dst_port_id = self.port_to_id.get(dst_port, 1)
                protocol = rule_state["protocol"]
                protocol_id = self.protocol_to_id.get(protocol, 1)
                obs[i] = {
                    "position": i - 1,
                    "permission": rule_state["action"],
                    "source_ip_id": src_node_id,
                    "source_wildcard_id": src_wildcard_id,
                    "source_port_id": src_port_id,
                    "dest_ip_id": dst_node_ip,
                    "dest_wildcard_id": dst_wildcard_id,
                    "dest_port_id": dst_port_id,
                    "protocol_id": protocol_id,
                }
            i += 1
        return obs

    @property
    def space(self) -> spaces.Space:
        return spaces.Dict(
            {
                i
                + 1: spaces.Dict(
                    {
                        "position": spaces.Discrete(self.num_rules),
                        "permission": spaces.Discrete(3),
                        # adding two to lengths is to account for reserved values 0 (unused) and 1 (any)
                        "source_ip_id": spaces.Discrete(len(self.ip_to_id) + 2),
                        "source_wildcard_id": spaces.Discrete(len(self.wildcard_to_id)+2),
                        "source_port_id": spaces.Discrete(len(self.port_to_id) + 2),
                        "dest_ip_id": spaces.Discrete(len(self.ip_to_id) + 2),
                        "dest_wildcard_id": spaces.Discrete(len(self.wildcard_to_id)+2),
                        "dest_port_id": spaces.Discrete(len(self.port_to_id) + 2),
                        "protocol_id": spaces.Discrete(len(self.protocol_to_id) + 2),
                    }
                )
                for i in range(self.num_rules)
            }
        )


    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ACLObservation:
        return cls(
            where = parent_where+["acl", "acl"],
            num_rules = config.num_rules,
            ip_list = config.ip_list,
            wildcard_list = config.wildcard_list,
            port_list = config.port_list,
            protocol_list = config.protocol_list
            )

class RouterObservation(AbstractObservation, identifier="ROUTER"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        hostname: str
        ports: Optional[List[PortObservation.ConfigSchema]] = None
        num_ports: Optional[int] = None
        acl: Optional[ACLObservation.ConfigSchema] = None
        ip_list: Optional[List[str]] = None
        wildcard_list: Optional[List[str]] = None
        port_list: Optional[List[int]] = None
        protocol_list: Optional[List[str]] = None
        num_rules: Optional[int] = None

    def __init__(self,
                 where: WhereType,
                 ports:List[PortObservation],
                 num_ports: int,
                 acl: ACLObservation,
                 )->None:
        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 = f"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) -> Any:
        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:
        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:
        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)

class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        hostname: str
        ip_list: Optional[List[str]] = None
        wildcard_list: Optional[List[str]] = None
        port_list: Optional[List[int]] = None
        protocol_list: Optional[List[str]] = None
        num_rules: Optional[int] = None


    def __init__(self,
                 where: WhereType,
                 ip_list: List[str],
                 wildcard_list: List[str],
                 port_list: List[int],
                 protocol_list: List[str],
                 num_rules: int,
                 )->None:
        self.where: WhereType = where

        self.ports: List[PortObservation] = [PortObservation(where=self.where+["port", port_num]) for port_num in (1,2,3) ]
        #TODO: check what the port nums are for firewall.

        self.internal_inbound_acl = ACLObservation(where = self.where+["acl","internal","inbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)
        self.internal_outbound_acl = ACLObservation(where = self.where+["acl","internal","outbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)
        self.dmz_inbound_acl = ACLObservation(where = self.where+["acl","dmz","inbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)
        self.dmz_outbound_acl = ACLObservation(where = self.where+["acl","dmz","outbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)
        self.external_inbound_acl = ACLObservation(where = self.where+["acl","external","inbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)
        self.external_outbound_acl = ACLObservation(where = self.where+["acl","external","outbound"], num_rules= num_rules, ip_list=ip_list, wildcard_list=wildcard_list, port_list=port_list, protocol_list=protocol_list)


        self.default_observation = {
            "PORTS": {i+1:p.default_observation for i,p in enumerate(self.ports)},
            "INTERNAL": {
                "INBOUND": self.internal_inbound_acl.default_observation,
                "OUTBOUND": self.internal_outbound_acl.default_observation,
                },
            "DMZ": {
                "INBOUND": self.dmz_inbound_acl.default_observation,
                "OUTBOUND": self.dmz_outbound_acl.default_observation,
                },
            "EXTERNAL": {
                "INBOUND": self.external_inbound_acl.default_observation,
                "OUTBOUND": self.external_outbound_acl.default_observation,
                },
            }

    def observe(self, state: Dict) -> Any:
        obs = {
            "PORTS": {i+1:p.observe(state) for i,p in enumerate(self.ports)},
            "INTERNAL": {
                "INBOUND": self.internal_inbound_acl.observe(state),
                "OUTBOUND": self.internal_outbound_acl.observe(state),
                },
            "DMZ": {
                "INBOUND": self.dmz_inbound_acl.observe(state),
                "OUTBOUND": self.dmz_outbound_acl.observe(state),
                },
            "EXTERNAL": {
                "INBOUND": self.external_inbound_acl.observe(state),
                "OUTBOUND": self.external_outbound_acl.observe(state),
                },
            }
        return obs

    @property
    def space(self) -> spaces.Space:
        space =spaces.Dict({
            "PORTS": spaces.Dict({i+1:p.space for i,p in enumerate(self.ports)}),
            "INTERNAL": spaces.Dict({
                "INBOUND": self.internal_inbound_acl.space,
                "OUTBOUND": self.internal_outbound_acl.space,
                }),
            "DMZ": spaces.Dict({
                "INBOUND": self.dmz_inbound_acl.space,
                "OUTBOUND": self.dmz_outbound_acl.space,
                }),
            "EXTERNAL": spaces.Dict({
                "INBOUND": self.external_inbound_acl.space,
                "OUTBOUND": self.external_outbound_acl.space,
                }),
            })
        return space

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> FirewallObservation:
        where = parent_where+["nodes", config.hostname]
        return cls(where=where, ip_list=config.ip_list, wildcard_list=config.wildcard_list, port_list=config.port_list, protocol_list=config.protocol_list, num_rules=config.num_rules)

class NodesObservation(AbstractObservation, identifier="NODES"):
    class ConfigSchema(AbstractObservation.ConfigSchema):
        """Config"""
        hosts: List[HostObservation.ConfigSchema] = []
        routers: List[RouterObservation.ConfigSchema] = []
        firewalls: List[FirewallObservation.ConfigSchema] = []

        num_services: int
        num_applications: int
        num_folders: int
        num_files: int
        num_nics: int
        include_nmne: bool
        include_num_access: bool

        num_ports: int
        ip_list: List[str]
        wildcard_list: List[str]
        port_list: List[int]
        protocol_list: List[str]
        num_rules: int


    def __init__(self, where: WhereType, hosts:List[HostObservation], routers:List[RouterObservation], firewalls:List[FirewallObservation])->None:
        self.where :WhereType = where

        self.hosts: List[HostObservation] = hosts
        self.routers: List[RouterObservation] = routers
        self.firewalls: List[FirewallObservation] = firewalls

        self.default_observation = {
            **{f"HOST{i}":host.default_observation for i,host in enumerate(self.hosts)},
            **{f"ROUTER{i}":router.default_observation for i,router in enumerate(self.routers)},
            **{f"FIREWALL{i}":firewall.default_observation for i,firewall in enumerate(self.firewalls)},
        }

    def observe(self, state: Dict) -> Any:
        obs = {
            **{f"HOST{i}":host.observe(state) for i,host in enumerate(self.hosts)},
            **{f"ROUTER{i}":router.observe(state) for i,router in enumerate(self.routers)},
            **{f"FIREWALL{i}":firewall.observe(state) for i,firewall in enumerate(self.firewalls)},
        }
        return obs

    @property
    def space(self) -> spaces.Space:
        space = spaces.Dict({
            **{f"HOST{i}":host.space for i,host in enumerate(self.hosts)},
            **{f"ROUTER{i}":router.space for i,router in enumerate(self.routers)},
            **{f"FIREWALL{i}":firewall.space for i,firewall in enumerate(self.firewalls)},
        })
        return space

    @classmethod
    def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
        if parent_where is None:
            where = ["network", "nodes"]
        else:
            where = parent_where + ["nodes"]

        for host_config in config.hosts:
            if host_config.num_services is None:
                host_config.num_services = config.num_services
            if host_config.num_applications is None:
                host_config.num_applications = config.num_applications
            if host_config.num_folders is None:
                host_config.num_folders = config.num_folders
            if host_config.num_files is None:
                host_config.num_files = config.num_files
            if host_config.num_nics is None:
                host_config.num_nics = config.num_nics
            if host_config.include_nmne is None:
                host_config.include_nmne = config.include_nmne
            if host_config.include_num_access is None:
                host_config.include_num_access = config.include_num_access

        for router_config in config.routers:
            if router_config.num_ports is None:
                router_config.num_ports = config.num_ports
            if router_config.ip_list is None:
                router_config.ip_list = config.ip_list
            if router_config.wildcard_list is None:
                router_config.wildcard_list = config.wildcard_list
            if router_config.port_list is None:
                router_config.port_list = config.port_list
            if router_config.protocol_list is None:
                router_config.protocol_list = config.protocol_list
            if router_config.num_rules is None:
                router_config.num_rules = config.num_rules

        for firewall_config in config.firewalls:
            if firewall_config.ip_list is None:
                firewall_config.ip_list = config.ip_list
            if firewall_config.wildcard_list is None:
                firewall_config.wildcard_list = config.wildcard_list
            if firewall_config.port_list is None:
                firewall_config.port_list = config.port_list
            if firewall_config.protocol_list is None:
                firewall_config.protocol_list = config.protocol_list
            if firewall_config.num_rules is None:
                firewall_config.num_rules = config.num_rules

        hosts = [HostObservation.from_config(config=c, parent_where=where) for c in config.hosts]
        routers = [RouterObservation.from_config(config=c, parent_where=where) for c in config.routers]
        firewalls = [FirewallObservation.from_config(config=c, parent_where=where) for c in config.firewalls]

        return cls(where=where, hosts=hosts, routers=routers, firewalls=firewalls)