From b1e46b4f9ebf3fcb84c0b96acbf4e3e9bdb1f689 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Thu, 14 Sep 2023 20:08:06 +0100 Subject: [PATCH] #1752: Apply suggestions from PR review --- .../system/applications/web_browser.py | 24 +------------ .../simulator/system/services/dns_client.py | 16 ++++----- .../simulator/system/services/dns_server.py | 34 +++---------------- 3 files changed, 14 insertions(+), 60 deletions(-) diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index b30f9946..78d196b7 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -1,6 +1,5 @@ -from abc import abstractmethod from ipaddress import IPv4Address -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional from primaite.simulator.system.applications.application import Application @@ -19,27 +18,6 @@ class WebBrowser(Application): history: Dict[str] "A dict that stores all of the previous domain names." - @abstractmethod - def describe_state(self) -> Dict: - """ - Describes the current state of the software. - - The specifics of the software's state, including its health, criticality, - and any other pertinent information, should be implemented in subclasses. - - :return: A dictionary containing key-value pairs representing the current state of the software. - :rtype: Dict - """ - pass - - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the Application. - - :param action: A list of actions to apply. - """ - pass - def reset_component_for_episode(self, episode: int): """ Resets the Application component for a new episode. diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index d6e4a05b..af20d6b8 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -1,5 +1,5 @@ from ipaddress import IPv4Address -from typing import Any, Dict, Optional +from typing import Dict, Optional from primaite import getLogger from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest @@ -36,7 +36,8 @@ class DNSClient(Service): :return: A dictionary containing key-value pairs representing the current state of the software. :rtype: Dict """ - return {"Operating State": self.operating_state} + state = super().describe_state() + return state def reset_component_for_episode(self, episode: int): """ @@ -45,8 +46,7 @@ class DNSClient(Service): This method ensures the Service is ready for a new episode, including resetting any stateful properties or statistics, and clearing any message queues. """ - super().reset_component_for_episode(episode=episode) - self.dns_cache = {} + pass def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address): """ @@ -68,8 +68,8 @@ class DNSClient(Service): """Function to check if domain name exists. :param: target_domain: The domain requested for an IP address. - :param: dest_ip_address: The ip address of the payload destination. - :param: dest_port: The port of the payload destination. + :param: dest_ip_address: The ip address of the DNS Server used for domain lookup. + :param: dest_port: The port on the DNS Server which accepts domain lookup requests. Default is Port.DNS. :param: session_id: The Session ID the payload is to originate from. Optional. :param: is_reattempt: Checks if the request has been reattempted. Default is False. """ @@ -105,7 +105,7 @@ class DNSClient(Service): def send( self, - payload: Any, + payload: DNSPacket, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -128,7 +128,7 @@ class DNSClient(Service): def receive( self, - payload: Any, + payload: DNSPacket, session_id: Optional[str] = None, **kwargs, ) -> bool: diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index c36c7034..be1145a2 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -37,9 +37,10 @@ class DNSServer(Service): :return: A dictionary containing key-value pairs representing the current state of the software. :rtype: Dict """ - return {"Operating State": self.operating_state} + state = super().describe_state() + return state - def dns_lookup(self, target_domain: Any) -> Optional[IPv4Address]: + def dns_lookup(self, target_domain: str) -> Optional[IPv4Address]: """ Attempts to find the IP address for a domain name. @@ -70,32 +71,7 @@ class DNSServer(Service): This method ensures the Service is ready for a new episode, including resetting any stateful properties or statistics, and clearing any message queues. """ - self.dns_table = {} - super().reset_component_for_episode(episode=episode) - - def send( - self, - payload: Any, - session_id: Optional[str] = None, - **kwargs, - ) -> bool: - """ - Sends a payload to 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 send. - :param: session_id: The id of the session - - :return: True if successful, False otherwise. - """ - try: - self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id) - return True - except Exception as e: - _LOGGER.error(e) - return False + pass def receive( self, @@ -110,7 +86,7 @@ class DNSServer(Service): is generated should be implemented in subclasses. :param: payload: The payload to send. - :param: session_id: The id of the session + :param: session_id: The id of the session. Optional. :return: True if DNS request returns a valid IP, otherwise, False """