#1752: Apply suggestions from PR review

This commit is contained in:
Czar.Echavez
2023-09-14 20:08:06 +01:00
parent 98e103a984
commit b1e46b4f9e
3 changed files with 14 additions and 60 deletions

View File

@@ -1,6 +1,5 @@
from abc import abstractmethod
from ipaddress import IPv4Address 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 from primaite.simulator.system.applications.application import Application
@@ -19,27 +18,6 @@ class WebBrowser(Application):
history: Dict[str] history: Dict[str]
"A dict that stores all of the previous domain names." "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): def reset_component_for_episode(self, episode: int):
""" """
Resets the Application component for a new episode. Resets the Application component for a new episode.

View File

@@ -1,5 +1,5 @@
from ipaddress import IPv4Address from ipaddress import IPv4Address
from typing import Any, Dict, Optional from typing import Dict, Optional
from primaite import getLogger from primaite import getLogger
from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest 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. :return: A dictionary containing key-value pairs representing the current state of the software.
:rtype: Dict :rtype: Dict
""" """
return {"Operating State": self.operating_state} state = super().describe_state()
return state
def reset_component_for_episode(self, episode: int): 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 This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues. stateful properties or statistics, and clearing any message queues.
""" """
super().reset_component_for_episode(episode=episode) pass
self.dns_cache = {}
def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address): 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. """Function to check if domain name exists.
:param: target_domain: The domain requested for an IP address. :param: target_domain: The domain requested for an IP address.
:param: dest_ip_address: The ip address of the payload destination. :param: dest_ip_address: The ip address of the DNS Server used for domain lookup.
:param: dest_port: The port of the payload destination. :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: 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. :param: is_reattempt: Checks if the request has been reattempted. Default is False.
""" """
@@ -105,7 +105,7 @@ class DNSClient(Service):
def send( def send(
self, self,
payload: Any, payload: DNSPacket,
session_id: Optional[str] = None, session_id: Optional[str] = None,
**kwargs, **kwargs,
) -> bool: ) -> bool:
@@ -128,7 +128,7 @@ class DNSClient(Service):
def receive( def receive(
self, self,
payload: Any, payload: DNSPacket,
session_id: Optional[str] = None, session_id: Optional[str] = None,
**kwargs, **kwargs,
) -> bool: ) -> bool:

View File

@@ -37,9 +37,10 @@ class DNSServer(Service):
:return: A dictionary containing key-value pairs representing the current state of the software. :return: A dictionary containing key-value pairs representing the current state of the software.
:rtype: Dict :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. 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 This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues. stateful properties or statistics, and clearing any message queues.
""" """
self.dns_table = {} pass
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
def receive( def receive(
self, self,
@@ -110,7 +86,7 @@ class DNSServer(Service):
is generated should be implemented in subclasses. is generated should be implemented in subclasses.
:param: payload: The payload to send. :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 :return: True if DNS request returns a valid IP, otherwise, False
""" """