PrimAITE/src/primaite/simulator/system/software.py

# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
"""Base class for system software."""
import copy
from abc import ABC, abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Union

from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, ConfigDict, Field

from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.core.sys_log import SysLog
from primaite.utils.validation.ip_protocol import IPProtocol, PROTOCOL_LOOKUP
from primaite.utils.validation.port import Port

if TYPE_CHECKING:
    from primaite.simulator.system.core.software_manager import SoftwareManager


class SoftwareType(Enum):
    """
    An enumeration representing the different types of software within a simulated environment.

    Members:
    - APPLICATION: User-facing programs that may perform input/output operations.
    - SERVICE: Represents programs that run in the background and may perform input/output operations.
    - PROCESS: Software executed by a Node that does not have the ability to performing input/output operations.
    """

    APPLICATION = 1
    "User-facing software that may perform input/output operations."
    SERVICE = 2
    "Software that runs in the background and may perform input/output operations."
    PROCESS = 3
    "Software executed by a Node that does not have the ability to performing input/output operations."


class SoftwareHealthState(Enum):
    """Enumeration of the Software Health States."""

    UNUSED = 0
    "Unused state."
    GOOD = 1
    "The software is in a good and healthy condition."
    FIXING = 2
    "The software is undergoing FIXING or updates."
    COMPROMISED = 3
    "The software's security has been compromised."
    OVERWHELMED = 4
    "he software is overwhelmed and not functioning properly."


class SoftwareCriticality(Enum):
    """Enumeration of Software Criticality Levels."""

    LOWEST = 1
    "The lowest level of criticality."
    LOW = 2
    "A low level of criticality."
    MEDIUM = 3
    "A medium level of criticality."
    HIGH = 4
    "A high level of criticality."
    HIGHEST = 5
    "The highest level of criticality."


class Software(SimComponent, ABC):
    """
    A base class representing software in a simulator environment.

    This class is intended to be subclassed by specific types of software entities.
    It outlines the fundamental attributes and behaviors expected of any software in the simulation.
    """

    class ConfigSchema(BaseModel, ABC):
        """Configurable options for all software."""

        model_config = ConfigDict(extra="forbid")
        starting_health_state: SoftwareHealthState = SoftwareHealthState.GOOD
        criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
        fixing_duration: int = 2

    config: ConfigSchema = Field(default_factory=lambda: Software.ConfigSchema())

    name: str
    "The name of the software."
    health_state_actual: SoftwareHealthState = SoftwareHealthState.UNUSED
    "The actual health state of the software."
    health_state_visible: SoftwareHealthState = SoftwareHealthState.UNUSED
    "The health state of the software visible to the red agent."
    fixing_count: int = 0
    "The count of patches applied to the software, defaults to 0."
    scanning_count: int = 0
    "The count of times the software has been scanned, defaults to 0."
    revealed_to_red: bool = False
    "Indicates if the software has been revealed to red agent, defaults is False."
    software_manager: Optional["SoftwareManager"] = None
    "An instance of Software Manager that is used by the parent node."
    sys_log: SysLog = None
    "An instance of SysLog that is used by the parent node."
    file_system: FileSystem
    "The FileSystem of the Node the Software is installed on."
    folder: Optional[Folder] = None
    "The folder on the file system the Software uses."
    _fixing_countdown: Optional[int] = None
    "Current number of ticks left to patch the software."

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.health_state_actual = self.config.starting_health_state  # don't remove this

    def _init_request_manager(self) -> RequestManager:
        """
        Initialise the request manager.

        More information in user guide and docstring for SimComponent._init_request_manager.
        """
        rm = super()._init_request_manager()
        rm.add_request(
            "compromise",
            RequestType(
                func=lambda request, context: RequestResponse.from_bool(
                    self.set_health_state(SoftwareHealthState.COMPROMISED)
                ),
            ),
        )
        rm.add_request(
            "fix",
            RequestType(
                func=lambda request, context: RequestResponse.from_bool(self.fix()),
            ),
        )
        rm.add_request("scan", RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan())))
        return rm

    def _get_session_details(self, session_id: str) -> Session:
        """
        Returns the Session object from the given session id.

        :param: session_id: ID of the session that needs details retrieved
        """
        return self.software_manager.session_manager.sessions_by_uuid[session_id]

    @abstractmethod
    def describe_state(self) -> Dict:
        """
        Produce a dictionary describing the current state of this object.

        Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.

        :return: Current state of this object and child objects.
        :rtype: Dict
        """
        state = super().describe_state()
        state.update(
            {
                "health_state_actual": self.health_state_actual.value,
                "health_state_visible": self.health_state_visible.value,
                "criticality": self.config.criticality.value,
                "fixing_count": self.fixing_count,
                "scanning_count": self.scanning_count,
                "revealed_to_red": self.revealed_to_red,
            }
        )
        return state

    def set_health_state(self, health_state: SoftwareHealthState) -> bool:
        """
        Assign a new health state to this software.

        Note: this should only be possible when the software is currently running, but the software base class has no
        operating state, only subclasses do. So subclasses will need to implement this check. TODO: check if this should
        be changed so that the base Software class has a running attr.

        :param health_state: New health state to assign to the software
        :type health_state: SoftwareHealthState
        """
        self.health_state_actual = health_state
        return True

    def install(self) -> None:
        """
        Perform first-time setup of this service on a node.

        This is an abstract class that should be overwritten by specific applications or services. It must be called
        after the service is already associate with a node. For example, a service may need to authenticate with a
        server during installation, or create files in the node's filesystem.
        """
        pass

    def uninstall(self) -> None:
        """Uninstall this service from a node.

        This is an abstract class that should be overwritten by applications or services. It must be called after the
        `install` method has already been run on that node. It should undo any installation steps, for example by
        deleting files, or contacting a server.
        """
        pass

    def scan(self) -> bool:
        """Update the observed health status to match the actual health status."""
        self.health_state_visible = self.health_state_actual
        return True

    def fix(self) -> bool:
        """Perform a fix on the software."""
        if self.health_state_actual in (SoftwareHealthState.COMPROMISED, SoftwareHealthState.GOOD):
            self._fixing_countdown = self.config.fixing_duration
            self.set_health_state(SoftwareHealthState.FIXING)
            return True
        return False

    def _update_fix_status(self) -> None:
        """Update the fix status of the software."""
        self._fixing_countdown -= 1
        if self._fixing_countdown <= 0:
            self.set_health_state(SoftwareHealthState.GOOD)
            self._fixing_countdown = None
            self.fixing_count += 1

    def reveal_to_red(self) -> None:
        """Reveals the software to the red agent."""
        self.revealed_to_red = True

    def apply_timestep(self, timestep: int) -> None:
        """
        Apply a single timestep to the software.

        :param timestep: The current timestep of the simulation.
        """
        super().apply_timestep(timestep)
        if self.health_state_actual == SoftwareHealthState.FIXING:
            self._update_fix_status()

    def pre_timestep(self, timestep: int) -> None:
        """Apply pre-timestep logic."""
        super().pre_timestep(timestep)


class IOSoftware(Software, ABC):
    """
    Represents software in a simulator environment that is capable of input/output operations.

    This base class is meant to be sub-classed by Application and Service classes. It provides the blueprint for
    Applications and Services that can receive payloads from a Node's SessionManager (corresponding to layer 5 in the
    OSI Model), process them according to their internals, and send a response payload back to the SessionManager if
    required.
    """

    class ConfigSchema(Software.ConfigSchema, ABC):
        """Configuration options for all IO Software."""

        listen_on_ports: Set[Port] = Field(default_factory=set)

    config: ConfigSchema = Field(default_factory=lambda: IOSoftware.ConfigSchema())

    installing_count: int = 0
    "The number of times the software has been installed. Default is 0."
    max_sessions: int = 100
    "The maximum number of sessions that the software can handle simultaneously. Default is 0."
    tcp: bool = True
    "Indicates if the software uses TCP protocol for communication. Default is True."
    udp: bool = True
    "Indicates if the software uses UDP protocol for communication. Default is True."
    port: Port
    "The port to which the software is connected."
    listen_on_ports: Set[Port] = Field(default_factory=set)
    "The set of ports to listen on."
    protocol: IPProtocol
    "The IP Protocol the Software operates on."
    _connections: Dict[str, Dict] = {}
    "Active connections."

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.listen_on_ports = self.config.listen_on_ports

    @abstractmethod
    def describe_state(self) -> Dict:
        """
        Produce a dictionary describing the current state of this object.

        Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.

        :return: Current state of this object and child objects.
        :rtype: Dict
        """
        state = super().describe_state()
        state.update(
            {
                "installing_count": self.installing_count,
                "max_sessions": self.max_sessions,
                "tcp": self.tcp,
                "udp": self.udp,
                "port": self.port,
            }
        )
        return state

    @abstractmethod
    def _can_perform_action(self) -> bool:
        """
        Checks if the software can perform actions.

        This is done by checking if the software is operating properly or the node it is installed
        in is operational.

        Returns true if the software can perform actions.
        """
        if self.software_manager and self.software_manager.node.operating_state != NodeOperatingState.ON:
            self.software_manager.node.sys_log.error(
                f"{self.name} Error: {self.software_manager.node.config.hostname} is not powered on."
            )
            return False
        return True

    @property
    def connections(self) -> Dict[str, Dict]:
        """Return the public version of connections."""
        return copy.copy(self._connections)

    def add_connection(self, connection_id: Union[str, int], session_id: Optional[str] = None) -> bool:
        """
        Create a new connection to this service.

        Returns true if connection successfully created

        :param: connection_id: UUID of the connection to create
        :type: string
        """
        # if over or at capacity, set to overwhelmed
        if len(self._connections) >= self.max_sessions:
            self.set_health_state(SoftwareHealthState.OVERWHELMED)
            self.sys_log.warning(f"{self.name}: Connection request ({connection_id}) declined. Service is at capacity.")
            return False
        else:
            # if service was previously overwhelmed, set to good because there is enough space for connections
            if self.health_state_actual == SoftwareHealthState.OVERWHELMED:
                self.set_health_state(SoftwareHealthState.GOOD)

            # check that connection already doesn't exist
            if not self._connections.get(connection_id):
                session_details = None
                if session_id:
                    session_details = self._get_session_details(session_id)
                self._connections[connection_id] = {
                    "session_id": session_id,
                    "ip_address": session_details.with_ip_address if session_details else None,
                    "time": datetime.now(),
                }
                self.sys_log.info(f"{self.name}: Connection request ({connection_id}) authorised")
                return True
            # connection with given id already exists
            self.sys_log.warning(
                f"{self.name}: Connection request ({connection_id}) declined. Connection already exists."
            )
            return False

    def terminate_connection(self, connection_id: str, send_disconnect: bool = True) -> bool:
        """
        Terminates a connection from this service.

        Returns true if connection successfully removed

        :param: connection_id: UUID of the connection to create
        :param send_disconnect: If True, sends a disconnect payload to the ip address of the associated session.
        :type: string
        """
        if self.connections.get(connection_id):
            connection_dict = self._connections.pop(connection_id)
            if send_disconnect:
                self.software_manager.send_payload_to_session_manager(
                    payload={"type": "disconnect", "connection_id": connection_id},
                    session_id=connection_dict["session_id"],
                )
                self.sys_log.info(f"{self.name}: Connection {connection_id=} terminated")
                return True
        return False

    def show_connections(self, markdown: bool = False):
        """
        Display the connections in tabular format.

        :param markdown: Whether to display the table in Markdown format or not. Default is `False`.
        """
        table = PrettyTable(["IP Address", "Connection ID", "Creation Timestamp"])
        if markdown:
            table.set_style(MARKDOWN)
        table.align = "l"
        table.title = f"{self.sys_log.hostname} {self.name} Connections"
        for connection_id, data in self.connections.items():
            table.add_row([data["ip_address"], connection_id, str(data["time"])])
        print(table.get_string(sortby="Creation Timestamp"))

    def clear_connections(self):
        """Clears all the connections from the software."""
        self._connections = {}

    def send(
        self,
        payload: Any,
        session_id: Optional[str] = None,
        dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
        dest_port: Optional[int] = None,
        ip_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"],
        **kwargs,
    ) -> bool:
        """
        Sends a payload to the SessionManager for network transmission.

        This method is responsible for initiating the process of sending network payloads. It supports both
        unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified
        as an IPv4Network. It delegates the actual sending process to the SoftwareManager.

        :param payload: The payload to be sent.
        :param dest_ip_address: The IP address or network (for broadcasts) of the payload destination.
        :param dest_port: The destination port for the payload. Optional.
        :param session_id: The Session ID from which the payload originates. Optional.
        :return: True if the payload was successfully sent, False otherwise.
        """
        if not self._can_perform_action():
            return False

        return self.software_manager.send_payload_to_session_manager(
            payload=payload,
            dest_ip_address=dest_ip_address,
            dest_port=dest_port,
            ip_protocol=ip_protocol,
            session_id=session_id,
        )

    @abstractmethod
    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
        """
        Receives a payload from the SessionManager.

        The specifics of how the payload is processed and whether a response payload
        is generated should be implemented in subclasses.


        :param payload: The payload to receive.
        :param session_id: The identifier of the session that the payload is associated with.
        :param kwargs: Additional keyword arguments specific to the implementation.
        :return: True if the payload was successfully received and processed, False otherwise.
        """
        # return false if not allowed to perform actions
        return self._can_perform_action()