diff --git a/src/primaite/simulator/network/protocols/dns.py b/src/primaite/simulator/network/protocols/dns.py index b8f0d8bd..0afa6405 100644 --- a/src/primaite/simulator/network/protocols/dns.py +++ b/src/primaite/simulator/network/protocols/dns.py @@ -5,38 +5,13 @@ from typing import Optional from pydantic import BaseModel -""" -class DNSEntry(BaseModel): - - Represents an entry in the DNS cache. - - :param domain_name: The domain name which a node would like to access. - :param ip_address: The IP address through which the domain name is reachable. - - - domain_name: str - ip_address: IPv4Address -""" - class DNSRequest(BaseModel): """Represents a DNS Request packet of a network frame. - :param sender_mac_addr: Sender MAC address. - :param sender_ip: Sender IP address. - :param target_mac_addr: Target MAC address. - :param target_ip: Target IP address. :param domain_name_request: Domain Name Request for IP address. """ - sender_mac_addr: str - "Sender MAC address." - sender_ip: IPv4Address - "Sender IP address." - target_mac_addr: Optional[str] = None - "Target MAC address of the DNS Server." - target_ip: IPv4Address - "Target IP address of the DNS Server." domain_name_request: str "Domain Name Request for IP address." @@ -44,21 +19,9 @@ class DNSRequest(BaseModel): class DNSReply(BaseModel): """Represents a DNS Reply packet of a network frame. - :param sender_mac_addr: Sender MAC address. - :param sender_ip: Sender IP address. - :param target_mac_addr: Target MAC address of DNS Client. - :param target_ip: Target IP address of DNS Client. :param domain_name_ip_address: IP Address of the Domain Name requested. """ - sender_mac_addr: str - "Sender MAC address." - sender_ip: IPv4Address - "Sender IP address." - target_mac_addr: Optional[str] = None - "Target MAC address of the DNS Server." - target_ip: IPv4Address - "Target IP address of the DNS Server." domain_name_ip_address: IPv4Address "IP Address of the Domain Name requested." @@ -73,15 +36,12 @@ class DNSPacket(BaseModel): :Example: >>> dns_request = DNSPacket( - ... dns_request=DNSRequest(sender_mac_addr="aa:bb:cc:dd:ee:ff", sender_ip = IPv4Address("192.168.0.1"), - ... target_ip = IPv4Address("192.168.0.2"), domain_name_request="www.google.co.uk"), + ... domain_name_request=DNSRequest(domain_name_request="www.google.co.uk"), ... dns_reply=None ... ) >>> dns_response = DNSPacket( - ... dns_request=DNSRequest(sender_mac_addr="aa:bb:cc:dd:ee:ff", sender_ip = IPv4Address("192.168.0.1"), - ... target_ip = IPv4Address("192.168.0.2"), domain_name_request="www.google.co.uk"), - ... dns_reply=DNSReply(sender_mac_addr="gg:hh:ii:jj:kk:ll", sender_ip = IPv4Address("192.168.0.2"), - ... target_ip = IPv4Address("192.168.0.1"), domain_name_ip_address=IPv4Address("142.250.179.227")) + ... dns_request=DNSRequest(domain_name_request="www.google.co.uk"), + ... dns_reply=DNSReply(domain_name_ip_address=IPv4Address("142.250.179.227")) ... ) """ @@ -97,18 +57,6 @@ class DNSPacket(BaseModel): :return: A new instance of DNSPacket. """ return DNSPacket( - dns_request=DNSRequest( - sender_mac_addr=self.dns_request.sender_mac_addr, - sender_ip=self.dns_request.sender_ip, - target_mac_addr=self.dns_request.target_mac_addr, - target_ip=self.dns_request.target_ip, - domain_name_request=self.dns_request.domain_name_request, - ), - dns_reply=DNSReply( - sender_mac_addr=self.dns_request.target_mac_addr, - sender_ip=self.dns_request.target_ip, - target_mac_addr=self.dns_request.sender_mac_addr, - target_ip=self.dns_request.sender_ip, - domain_name_ip_address=domain_ip_address, - ), + dns_request=DNSRequest(domain_name_request=self.dns_request.domain_name_request), + dns_reply=DNSReply(domain_name_ip_address=domain_ip_address), ) diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index 1b7ccf7d..e01b8d2e 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -5,6 +5,7 @@ from pydantic import BaseModel from primaite import getLogger from primaite.simulator.network.protocols.arp import ARPPacket +from primaite.simulator.network.protocols.dns import DNSPacket from primaite.simulator.network.transmission.network_layer import ICMPPacket, IPPacket, IPProtocol from primaite.simulator.network.transmission.primaite_layer import PrimaiteHeader from primaite.simulator.network.transmission.transport_layer import TCPHeader, UDPHeader @@ -96,6 +97,8 @@ class Frame(BaseModel): "ICMP header." arp: Optional[ARPPacket] = None "ARP packet." + dns: Optional[DNSPacket] = None + "DNS packet." primaite: PrimaiteHeader "PrimAITE header." payload: Optional[Any] = None diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index 7b080244..ce4a9150 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -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 diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index fe9a6123..7ad3bac1 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -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