#1752 - Changed DNSReply and DNSResponse to have 1 parameter only

This commit is contained in:
SunilSamra
2023-08-21 14:11:53 +01:00
parent a0b258a597
commit 1a13af2f5e
4 changed files with 41 additions and 67 deletions

View File

@@ -4,14 +4,14 @@ from typing import Any, Dict, List
from pydantic import BaseModel
from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest
class DNSClient(BaseModel):
"""Represents a DNS Client as a Service."""
target_url: str
"The URL/domain name the client requests the IP for."
dns_cache: Dict[str:IPv4Address] = {}
"A dict of known mappings between domain names and IPv4 addresses."
"A dict of known mappings between domain/URLs names and IPv4 addresses."
@abstractmethod
def describe_state(self) -> Dict:
@@ -41,9 +41,16 @@ class DNSClient(BaseModel):
This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues.
"""
self.target_url = ""
self.dns_cache = {}
def check_domain_is_in_cache(self, target_domain: str, session_id: str):
"""Function to check if domain name is in DNS client cache."""
if target_domain in self.dns_cache:
ip_address = self.dns_cache[target_domain]
self.send(ip_address, session_id)
else:
self.send(target_domain, session_id)
def send(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
Sends a payload to the SessionManager.
@@ -54,7 +61,7 @@ class DNSClient(BaseModel):
:param payload: The payload to send.
:return: True if successful, False otherwise.
"""
pass
DNSPacket(dns_request=DNSRequest(domain_name_request=payload), dns_reply=None)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
@@ -63,7 +70,10 @@ class DNSClient(BaseModel):
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 payload: The payload to receive. (receive a DNS packet with dns request and dns reply in, send to web
browser)
:return: True if successful, False otherwise.
"""
# check DNS packet (dns request, dns reply) here and see if it actually worked
pass

View File

@@ -4,6 +4,8 @@ from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from primaite.simulator.network.protocols.dns import DNSPacket, DNSReply, DNSRequest
class DNSServer(BaseModel):
"""Represents a DNS Server as a Service."""
@@ -28,7 +30,7 @@ class DNSServer(BaseModel):
"""
Applies a list of actions to the Service.
:param action: A list of actions to apply.
:param action: A list of actions to apply. (unsure)
"""
pass
@@ -63,7 +65,13 @@ class DNSServer(BaseModel):
:param payload: The payload to send.
:return: True if successful, False otherwise.
"""
pass
ip_addresses = list(self.dns_table.values())
domain_names = list(self.dns_table.keys())
index_of_domain = ip_addresses.index(payload)
DNSPacket(
dns_request=DNSRequest(domain_name_request=domain_names[index_of_domain]),
dns_reply=DNSReply(domain_name_ip_address=payload),
)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
@@ -72,7 +80,12 @@ class DNSServer(BaseModel):
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 payload: The payload to receive. (take the domain name and do dns lookup)
:return: True if successful, False otherwise.
"""
pass
ip_address = self.dns_lookup(payload)
if ip_address is not None:
self.send(ip_address, session_id)
return True
return False