From 76d404dc2f4b954805d8970e9387eac49abc479a Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 15 Aug 2023 13:56:38 +0100 Subject: [PATCH 01/16] #1752 - Added dns_client.py and dns_server.py service files - Added new get_install method to software.py --- .../simulator/system/services/dns_client.py | 68 +++++++++++++++++++ .../simulator/system/services/dns_server.py | 64 +++++++++++++++++ src/primaite/simulator/system/software.py | 9 +++ 3 files changed, 141 insertions(+) create mode 100644 src/primaite/simulator/system/services/dns_client.py create mode 100644 src/primaite/simulator/system/services/dns_server.py diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py new file mode 100644 index 00000000..ebbe4058 --- /dev/null +++ b/src/primaite/simulator/system/services/dns_client.py @@ -0,0 +1,68 @@ +from abc import abstractmethod +from ipaddress import IPv4Address +from typing import Any, Dict, List + +from pydantic import BaseModel + + +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." + + @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 + """ + return {"Operating State": self.operating_state} + + def apply_action(self, action: List[str]) -> None: + """ + Applies a list of actions to the Service. + + :param action: A list of actions to apply. + """ + pass + + def reset_component_for_episode(self): + """ + Resets the Service component for a new episode. + + This method ensures the Service is ready for a new episode, including resetting any + stateful properties or statistics, and clearing any message queues. + """ + pass + + def send(self, payload: Any, session_id: str, **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. + :return: True if successful, False otherwise. + """ + pass + + 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. + :return: True if successful, False otherwise. + """ + pass diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py new file mode 100644 index 00000000..e7b51f38 --- /dev/null +++ b/src/primaite/simulator/system/services/dns_server.py @@ -0,0 +1,64 @@ +from abc import abstractmethod +from typing import Any, Dict, List + +from pydantic import BaseModel + + +class DNSServer(BaseModel): + """Represents a DNS Server as a Service.""" + + dns_table: dict[str:str] = {"https://google.co.uk": "8.8.8.8"} + + @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 + """ + return {"Operating State": self.operating_state} + + def apply_action(self, action: List[str]) -> None: + """ + Applies a list of actions to the Service. + + :param action: A list of actions to apply. + """ + pass + + def reset_component_for_episode(self): + """ + Resets the Service component for a new episode. + + This method ensures the Service is ready for a new episode, including resetting any + stateful properties or statistics, and clearing any message queues. + """ + pass + + def send(self, payload: Any, session_id: str, **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. + :return: True if successful, False otherwise. + """ + pass + + 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. + :return: True if successful, False otherwise. + """ + pass diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 854e7e2b..4caf6f03 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -109,6 +109,15 @@ class Software(SimComponent): """ pass + @staticmethod + def get_install(): + """ + This method ensures the software has to have a way to install it. + + This can be used by the software manager to install the software. + """ + pass + class IOSoftware(Software): """ From 72cd9fd8e29ed1d6c917f9bc9fad1d2fda5d0da5 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Wed, 16 Aug 2023 13:00:16 +0100 Subject: [PATCH 02/16] #1752 - Created dns.py protocol file with DNSPacket and DNSRequest and DNSReply packets - Added reset_component logic for dns_server.py and dns_client.py --- .../simulator/network/protocols/dns.py | 113 ++++++++++++++++++ .../simulator/system/services/dns_client.py | 3 +- .../simulator/system/services/dns_server.py | 4 +- 3 files changed, 117 insertions(+), 3 deletions(-) create mode 100644 src/primaite/simulator/network/protocols/dns.py diff --git a/src/primaite/simulator/network/protocols/dns.py b/src/primaite/simulator/network/protocols/dns.py new file mode 100644 index 00000000..fa88652c --- /dev/null +++ b/src/primaite/simulator/network/protocols/dns.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from ipaddress import IPv4Address +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." + + +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." + + +class DNSPacket(BaseModel): + """ + Represents the DNS layer of a network frame. + + :param dns_request: DNS Request packet sent by DNS Client. + :param dns_reply: DNS Reply packet generated by DNS Server. + + :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"), + ... 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 + "DNS Request packet sent by DNS Client." + dns_reply: Optional[DNSReply] = None + "DNS Reply packet generated by DNS Server." + + def generate_reply(self, domain_ip_address: IPv4Address) -> DNSPacket: + """Generate a new DNSPacket to be sent as a response with a DNS Reply packet which contains the IP address. + + :param domain_ip_address: The IP address that was being sought after from the original target domain name. + :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, + ), + ) diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index ebbe4058..7b080244 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -41,7 +41,8 @@ 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. """ - pass + self.target_url = "" + self.dns_cache = {} def send(self, payload: Any, session_id: str, **kwargs) -> bool: """ diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index e7b51f38..7eed51b5 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -7,7 +7,7 @@ from pydantic import BaseModel class DNSServer(BaseModel): """Represents a DNS Server as a Service.""" - dns_table: dict[str:str] = {"https://google.co.uk": "8.8.8.8"} + dns_table: dict[str:str] = {} @abstractmethod def describe_state(self) -> Dict: @@ -37,7 +37,7 @@ class DNSServer(BaseModel): This method ensures the Service is ready for a new episode, including resetting any stateful properties or statistics, and clearing any message queues. """ - pass + self.dns_table = {} def send(self, payload: Any, session_id: str, **kwargs) -> bool: """ From 2919be3796e41e1694c1e0e5cdd67b237f529ad7 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Thu, 17 Aug 2023 14:20:09 +0100 Subject: [PATCH 03/16] #1752 - Added web_browser.py application for DNS modelling --- .../simulator/network/protocols/dns.py | 7 +- .../system/applications/web_browser.py | 76 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) create mode 100644 src/primaite/simulator/system/applications/web_browser.py diff --git a/src/primaite/simulator/network/protocols/dns.py b/src/primaite/simulator/network/protocols/dns.py index fa88652c..b8f0d8bd 100644 --- a/src/primaite/simulator/network/protocols/dns.py +++ b/src/primaite/simulator/network/protocols/dns.py @@ -5,17 +5,18 @@ 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): diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py new file mode 100644 index 00000000..b30f9946 --- /dev/null +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -0,0 +1,76 @@ +from abc import abstractmethod +from ipaddress import IPv4Address +from typing import Any, Dict, List, Optional + +from primaite.simulator.system.applications.application import Application + + +class WebBrowser(Application): + """ + Represents a web browser in the simulation environment. + + The application requests and loads web pages using its domain name and requesting IP addresses using DNS. + """ + + domain_name: str + "The domain name of the webpage." + domain_name_ip_address: Optional[IPv4Address] + "The IP address of the domain name for the webpage." + 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. + + This method ensures the Application is ready for a new episode, including resetting any + stateful properties or statistics, and clearing any message queues. + """ + self.domain_name = "" + self.domain_name_ip_address = None + self.history = {} + + def send(self, payload: Any, session_id: str, **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. + :return: True if successful, False otherwise. + """ + pass + + 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. + :return: True if successful, False otherwise. + """ + pass From a0b258a597bbcc71fe9f2c689a6a268d3de1aeff Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 21 Aug 2023 09:02:04 +0100 Subject: [PATCH 04/16] #1752 - Added a dns_lookup function to dns_server.py --- .../simulator/system/services/dns_server.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index 7eed51b5..fe9a6123 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -1,5 +1,6 @@ from abc import abstractmethod -from typing import Any, Dict, List +from ipaddress import IPv4Address +from typing import Any, Dict, List, Optional from pydantic import BaseModel @@ -7,7 +8,8 @@ from pydantic import BaseModel class DNSServer(BaseModel): """Represents a DNS Server as a Service.""" - dns_table: dict[str:str] = {} + dns_table: dict[str:IPv4Address] = {} + "A dict of mappings between domain names and IPv4 addresses." @abstractmethod def describe_state(self) -> Dict: @@ -30,6 +32,18 @@ class DNSServer(BaseModel): """ pass + def dns_lookup(self, target_domain: str) -> Optional[IPv4Address]: + """ + Attempts to find the IP address for a domain name. + + :param target_domain: The single domain name requested by a DNS client. + :return ip_address: The IP address of that domain name or None. + """ + if target_domain in self.dns_table: + return self.dns_table[target_domain] + else: + return None + def reset_component_for_episode(self): """ Resets the Service component for a new episode. From 1a13af2f5ec563b5060478bc00c9148105fe5fff Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 21 Aug 2023 14:11:53 +0100 Subject: [PATCH 05/16] #1752 - Changed DNSReply and DNSResponse to have 1 parameter only --- .../simulator/network/protocols/dns.py | 62 ++----------------- .../network/transmission/data_link_layer.py | 3 + .../simulator/system/services/dns_client.py | 22 +++++-- .../simulator/system/services/dns_server.py | 21 +++++-- 4 files changed, 41 insertions(+), 67 deletions(-) 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 From 550b62f75dc7b5bf29d43808feaf2b62687f0010 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Mon, 21 Aug 2023 16:09:17 +0100 Subject: [PATCH 06/16] #1752 - Removed unnecessary print statement - Changed docstring on function check_domain_in_cache --- src/primaite/simulator/system/services/dns_client.py | 8 ++++++-- tests/test_seeding_and_deterministic_session.py | 1 - 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index ce4a9150..f99411d2 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -43,8 +43,12 @@ class DNSClient(BaseModel): """ 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.""" + def check_domain_in_cache(self, target_domain: str, session_id: str): + """Function to check if domain name is in DNS client cache. + + :param target_domain: The domain requested for an IP address. + :param session_id: The ID of the session in order to send the response to the DNS server or application. + """ if target_domain in self.dns_cache: ip_address = self.dns_cache[target_domain] self.send(ip_address, session_id) diff --git a/tests/test_seeding_and_deterministic_session.py b/tests/test_seeding_and_deterministic_session.py index 9500c4a3..aff5496a 100644 --- a/tests/test_seeding_and_deterministic_session.py +++ b/tests/test_seeding_and_deterministic_session.py @@ -45,7 +45,6 @@ def test_seeded_learning(temp_primaite_session): ), "Expected output is based upon a agent that was trained with seed 67890" session.learn() actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict() - print(actual_mean_reward_per_episode, "THISt") assert actual_mean_reward_per_episode == expected_mean_reward_per_episode From c1ba3b0850e35306b80ee48de430ff472a316caa Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 22 Aug 2023 21:43:57 +0100 Subject: [PATCH 07/16] #1752 - Added comments to ticket --- src/primaite/simulator/system/services/dns_client.py | 3 +-- src/primaite/simulator/system/services/dns_server.py | 8 +++----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index f99411d2..97968407 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -78,6 +78,5 @@ class DNSClient(BaseModel): browser) :return: True if successful, False otherwise. """ - # check DNS packet (dns request, dns reply) here and see if it actually worked - + # check the 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 7ad3bac1..a2eaf9d9 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -42,7 +42,7 @@ class DNSServer(BaseModel): :return ip_address: The IP address of that domain name or None. """ if target_domain in self.dns_table: - return self.dns_table[target_domain] + self.dns_table[target_domain] else: return None @@ -65,11 +65,9 @@ class DNSServer(BaseModel): :param payload: The payload to send. :return: True if successful, False otherwise. """ - ip_addresses = list(self.dns_table.values()) - domain_names = list(self.dns_table.keys()) - index_of_domain = ip_addresses.index(payload) + # DNS packet will be sent from DNS Server to the DNS client DNSPacket( - dns_request=DNSRequest(domain_name_request=domain_names[index_of_domain]), + dns_request=DNSRequest(domain_name_request=self.dns_table), dns_reply=DNSReply(domain_name_ip_address=payload), ) From 47dd23311bd951a81e867a5e28186840c891dea9 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 7 Sep 2023 15:45:37 +0100 Subject: [PATCH 08/16] #1752: added more functionality to DNS client and server + tests --- src/primaite/simulator/core.py | 3 + .../simulator/system/core/session_manager.py | 81 +++++------- .../simulator/system/core/software_manager.py | 18 ++- .../simulator/system/services/dns_client.py | 119 +++++++++++++----- .../simulator/system/services/dns_server.py | 118 ++++++++++++----- .../simulator/system/services/service.py | 39 +++++- .../_simulator/_system/_services/test_dns.py | 66 ++++++++++ 7 files changed, 321 insertions(+), 123 deletions(-) create mode 100644 tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index 32db95c6..ee19abb3 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -192,6 +192,9 @@ class SimComponent(BaseModel): :param action: List describing the action to apply to this object. :type action: List[str] + + :param: context: Dict containing context for actions + :type context: Dict """ if self.action_manager is None: return diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index be20a28d..f8e97442 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -32,27 +32,23 @@ class Session(SimComponent): """ protocol: IPProtocol - src_ip_address: IPv4Address - dst_ip_address: IPv4Address + with_ip_address: IPv4Address src_port: Optional[Port] dst_port: Optional[Port] connected: bool = False @classmethod - def from_session_key( - cls, session_key: Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]] - ) -> Session: + def from_session_key(cls, session_key: Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]) -> Session: """ Create a Session instance from a session key tuple. :param session_key: Tuple containing the session details. :return: A Session instance. """ - protocol, src_ip_address, dst_ip_address, src_port, dst_port = session_key + protocol, with_ip_address, src_port, dst_port = session_key return Session( protocol=protocol, - src_ip_address=src_ip_address, - dst_ip_address=dst_ip_address, + with_ip_address=with_ip_address, src_port=src_port, dst_port=dst_port, ) @@ -78,9 +74,7 @@ class SessionManager: """ def __init__(self, sys_log: SysLog, arp_cache: "ARPCache"): - self.sessions_by_key: Dict[ - Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]], Session - ] = {} + self.sessions_by_key: Dict[Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]], Session] = {} self.sessions_by_uuid: Dict[str, Session] = {} self.sys_log: SysLog = sys_log self.software_manager: SoftwareManager = None # Noqa @@ -99,8 +93,8 @@ class SessionManager: @staticmethod def _get_session_key( - frame: Frame, from_source: bool = True - ) -> Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]]: + frame: Frame, inbound_frame: bool = True + ) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]: """ Extracts the session key from the given frame. @@ -112,36 +106,36 @@ class SessionManager: - Optional[Port]: The destination port number (if applicable). :param frame: The network frame from which to extract the session key. - :param from_source: A flag to indicate if the key should be extracted from the source or destination. :return: A tuple containing the session key. """ protocol = frame.ip.protocol - src_ip_address = frame.ip.src_ip_address - dst_ip_address = frame.ip.dst_ip_address + with_ip_address = frame.ip.src_ip_address if protocol == IPProtocol.TCP: - if from_source: + if inbound_frame: src_port = frame.tcp.src_port dst_port = frame.tcp.dst_port else: dst_port = frame.tcp.src_port src_port = frame.tcp.dst_port + with_ip_address = frame.ip.dst_ip_address elif protocol == IPProtocol.UDP: - if from_source: + if inbound_frame: src_port = frame.udp.src_port dst_port = frame.udp.dst_port else: dst_port = frame.udp.src_port src_port = frame.udp.dst_port + with_ip_address = frame.ip.dst_ip_address else: src_port = None dst_port = None - return protocol, src_ip_address, dst_ip_address, src_port, dst_port + return protocol, with_ip_address, src_port, dst_port def receive_payload_from_software_manager( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, + dst_ip_address: Optional[IPv4Address] = None, + dst_port: Optional[Port] = None, session_id: Optional[str] = None, is_reattempt: bool = False, ) -> Union[Any, None]: @@ -154,20 +148,21 @@ class SessionManager: :param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created. """ if session_id: - dest_ip_address = self.sessions_by_uuid[session_id].dst_ip_address - dest_port = self.sessions_by_uuid[session_id].dst_port + session = self.sessions_by_uuid[session_id] + dst_ip_address = self.sessions_by_uuid[session_id].with_ip_address + dst_port = self.sessions_by_uuid[session_id].dst_port - dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dest_ip_address) + dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address) if dst_mac_address: - outbound_nic = self.arp_cache.get_arp_cache_nic(dest_ip_address) + outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address) else: if not is_reattempt: - self.arp_cache.send_arp_request(dest_ip_address) + self.arp_cache.send_arp_request(dst_ip_address) return self.receive_payload_from_software_manager( payload=payload, - dest_ip_address=dest_ip_address, - dest_port=dest_port, + dst_ip_address=dst_ip_address, + dst_port=dst_port, session_id=session_id, is_reattempt=True, ) @@ -178,17 +173,17 @@ class SessionManager: ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address), ip=IPPacket( src_ip_address=outbound_nic.ip_address, - dst_ip_address=dest_ip_address, + dst_ip_address=dst_ip_address, ), tcp=TCPHeader( - src_port=dest_port, - dst_port=dest_port, + src_port=dst_port, + dst_port=dst_port, ), payload=payload, ) if not session_id: - session_key = self._get_session_key(frame, from_source=True) + session_key = self._get_session_key(frame, inbound_frame=False) session = self.sessions_by_key.get(session_key) if not session: # Create new session @@ -198,33 +193,25 @@ class SessionManager: outbound_nic.send_frame(frame) - def send_payload_to_software_manager(self, payload: Any, session_id: int): + def receive_frame(self, frame: Frame): """ - Send a payload to the software manager. - - :param payload: The payload to be sent. - :param session_id: The Session ID the payload originates from. - """ - self.software_manager.receive_payload_from_session_manger() - - def receive_payload_from_nic(self, frame: Frame): - """ - Receive a Frame from the NIC. + Receive a Frame. Extract the session key using the _get_session_key method, and forward the payload to the appropriate session. If the session does not exist, a new one is created. :param frame: The frame being received. """ - session_key = self._get_session_key(frame) - session = self.sessions_by_key.get(session_key) + session_key = self._get_session_key(frame, inbound_frame=True) + session: Session = self.sessions_by_key.get(session_key) if not session: # Create new session session = Session.from_session_key(session_key) self.sessions_by_key[session_key] = session self.sessions_by_uuid[session.uuid] = session - self.software_manager.receive_payload_from_session_manger(payload=frame, session=session) - # TODO: Implement the frame deconstruction and send to SoftwareManager. + self.software_manager.receive_payload_from_session_manager( + payload=frame.payload, port=frame.tcp.dst_port, protocol=frame.ip.protocol, session_id=session.uuid + ) def show(self, markdown: bool = False): """ diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 28e37963..71519ac7 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -6,7 +6,6 @@ from prettytable import MARKDOWN, PrettyTable from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.application import Application -from primaite.simulator.system.core.session_manager import Session from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.services.service import Service from primaite.simulator.system.software import SoftwareType @@ -86,7 +85,7 @@ class SoftwareManager: payload: Any, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = None, - session_id: Optional[int] = None, + session_id: Optional[str] = None, ): """ Send a payload to the SessionManager. @@ -97,22 +96,21 @@ class SoftwareManager: :param session_id: The Session ID the payload is to originate from. Optional. """ self.session_manager.receive_payload_from_software_manager( - payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id + payload=payload, dst_ip_address=dest_ip_address, dst_port=dest_port, session_id=session_id ) - def receive_payload_from_session_manger(self, payload: Any, session: Session): + def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str): """ Receive a payload from the SessionManager and forward it to the corresponding service or application. :param payload: The payload being received. :param session: The transport session the payload originates from. """ - # receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None) - # if receiver: - # receiver.receive_payload(None, payload) - # else: - # raise ValueError(f"No service or application found for port {port} and protocol {protocol}") - pass + receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None) + if receiver: + receiver.receive_payload(payload=payload, session_id=session_id) + else: + self.sys_log.error(f"No service or application found for port {port} and protocol {protocol}") def show(self, markdown: bool = False): """ diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index 97968407..3929065d 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -1,19 +1,26 @@ -from abc import abstractmethod from ipaddress import IPv4Address -from typing import Any, Dict, List - -from pydantic import BaseModel +from typing import Any, Dict, Optional from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.services.service import Service -class DNSClient(BaseModel): +class DNSClient(Service): """Represents a DNS Client as a Service.""" - dns_cache: Dict[str:IPv4Address] = {} + dns_cache: Dict[str, IPv4Address] = {} "A dict of known mappings between domain/URLs names and IPv4 addresses." - @abstractmethod + def __init__(self, **kwargs): + kwargs["name"] = "DNSClient" + kwargs["port"] = Port.DNS + # DNS uses UDP by default + # it switches to TCP when the bytes exceed 512 (or 4096) bytes + kwargs["protocol"] = IPProtocol.UDP + super().__init__(**kwargs) + def describe_state(self) -> Dict: """ Describes the current state of the software. @@ -26,57 +33,109 @@ class DNSClient(BaseModel): """ return {"Operating State": self.operating_state} - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the Service. - - :param action: A list of actions to apply. - """ - pass - - def reset_component_for_episode(self): + def reset_component_for_episode(self, episode: int): """ Resets the Service component for a new episode. 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 = {} - def check_domain_in_cache(self, target_domain: str, session_id: str): + def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address): + """ + Adds a domain name to the DNS Client cache. + + :param: domain_name: The domain name to save to cache + :param: ip_address: The IP Address to attach the domain name to + """ + self.dns_cache[domain_name] = ip_address + + def check_domain_in_cache( + self, + target_domain: str, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + session_id: Optional[str] = None, + is_reattempt: bool = False, + ) -> bool: """Function to check if domain name is in DNS client cache. - :param target_domain: The domain requested for an IP address. - :param session_id: The ID of the session in order to send the response to the DNS server or application. + :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: 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. """ - 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) + # check if the target domain is in the client's DNS cache + payload = DNSPacket(dns_request=DNSRequest(domain_name_request=target_domain)) - def send(self, payload: Any, session_id: str, **kwargs) -> bool: + # check if the domain is already in the DNS cache + if target_domain in self.dns_cache: + return True + else: + # return False if already reattempted + if is_reattempt: + return False + else: + # send a request to check if domain name exists in the DNS Server + self.send(payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id) + # call function again + return self.check_domain_in_cache( + target_domain=target_domain, + dest_ip_address=dest_ip_address, + dest_port=dest_port, + session_id=session_id, + is_reattempt=True, + ) + + def send( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + 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 payload: The payload to be sent. + :param dest_ip_address: The ip address of the payload destination. + :param dest_port: The port of the payload destination. + :param session_id: The Session ID the payload is to originate from. Optional. + :return: True if successful, False otherwise. """ - DNSPacket(dns_request=DNSRequest(domain_name_request=payload), dns_reply=None) + # create DNS request packet + self.software_manager.send_payload_to_session_manager( + payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id + ) - def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + def receive( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + session_id: Optional[str] = None, + **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. (receive a DNS packet with dns request and dns reply in, send to web - browser) + :param payload: The payload to be sent. + :param dest_ip_address: The ip address of the payload destination. + :param dest_port: The port of the payload destination. + :param session_id: The Session ID the payload is to originate from. Optional. :return: True if successful, False otherwise. """ + super().send() # check the 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 a2eaf9d9..3dcd89f9 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -1,19 +1,31 @@ -from abc import abstractmethod from ipaddress import IPv4Address -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional -from pydantic import BaseModel +from prettytable import MARKDOWN, PrettyTable -from primaite.simulator.network.protocols.dns import DNSPacket, DNSReply, DNSRequest +from primaite import getLogger +from primaite.simulator.network.protocols.dns import DNSPacket +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.services.service import Service + +_LOGGER = getLogger(__name__) -class DNSServer(BaseModel): +class DNSServer(Service): """Represents a DNS Server as a Service.""" - dns_table: dict[str:IPv4Address] = {} + dns_table: Dict[str, IPv4Address] = {} "A dict of mappings between domain names and IPv4 addresses." - @abstractmethod + def __init__(self, **kwargs): + kwargs["name"] = "DNSServer" + kwargs["port"] = Port.DNS + # DNS uses UDP by default + # it switches to TCP when the bytes exceed 512 (or 4096) bytes + kwargs["protocol"] = IPProtocol.UDP + super().__init__(**kwargs) + def describe_state(self) -> Dict: """ Describes the current state of the software. @@ -26,15 +38,7 @@ class DNSServer(BaseModel): """ return {"Operating State": self.operating_state} - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the Service. - - :param action: A list of actions to apply. (unsure) - """ - pass - - def dns_lookup(self, target_domain: str) -> Optional[IPv4Address]: + def dns_lookup(self, target_domain: Any) -> Optional[IPv4Address]: """ Attempts to find the IP address for a domain name. @@ -42,11 +46,23 @@ class DNSServer(BaseModel): :return ip_address: The IP address of that domain name or None. """ if target_domain in self.dns_table: - self.dns_table[target_domain] + return self.dns_table[target_domain] else: return None - def reset_component_for_episode(self): + def dns_register(self, domain_name: str, domain_ip_address: IPv4Address): + """ + Register a domain name and its IP address. + + :param: domain_name: The domain name to register + :type: domain_name: str + + :param: domain_ip_address: The IP address that the domain should route to + :type: domain_ip_address: IPv4Address + """ + self.dns_table[domain_name] = domain_ip_address + + def reset_component_for_episode(self, episode: int): """ Resets the Service component for a new episode. @@ -54,36 +70,78 @@ class DNSServer(BaseModel): 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: str, **kwargs) -> bool: + def send( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + 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: payload: The payload to send. + :param: dest_ip_address: The ip address of the machine that the payload will be sent to + :param: dest_port: The port of the machine that the payload will be sent to + :param: session_id: The id of the session + :return: True if successful, False otherwise. """ - # DNS packet will be sent from DNS Server to the DNS client - DNSPacket( - dns_request=DNSRequest(domain_name_request=self.dns_table), - dns_reply=DNSReply(domain_name_ip_address=payload), - ) + try: + self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id) + except Exception as e: + _LOGGER.error(e) + return False - def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + def receive( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + session_id: Optional[str] = None, + **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. (take the domain name and do dns lookup) + :param: payload: The payload to send. + :param: dest_ip_address: The ip address of the machine that the payload will be sent to + :param: dest_port: The port of the machine that the payload will be sent to + :param: session_id: The id of the session + :return: True if successful, False otherwise. """ - ip_address = self.dns_lookup(payload) - if ip_address is not None: - self.send(ip_address, session_id) + # The payload should be a DNS packet + if not isinstance(payload, DNSPacket): + _LOGGER.debug(f"{payload} is not a DNSPacket") + return False + # cast payload into a DNS packet + payload: DNSPacket = payload + if payload.dns_request is not None: + # generate a reply with the correct DNS IP address + payload.generate_reply(self.dns_lookup(payload.dns_request.domain_name_request)) + # send reply + self.send(payload, session_id) return True return False + + def show(self, markdown: bool = False): + """Prints a table of DNS Lookup table.""" + table = PrettyTable(["Domain Name", "IP Address"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"{self.sys_log.hostname} DNS Lookup table" + for dns in self.dns_table.items(): + table.add_row([dns[0], dns[1]]) + print(table) diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index b9340103..3011c74d 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -1,8 +1,10 @@ from enum import Enum +from ipaddress import IPv4Address from typing import Any, Dict, Optional from primaite import getLogger from primaite.simulator.core import Action, ActionManager +from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.software import IOSoftware _LOGGER = getLogger(__name__) @@ -72,29 +74,54 @@ class Service(IOSoftware): """ pass - def send(self, payload: Any, session_id: str, **kwargs) -> bool: + def send( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + 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: payload: The payload to send. + :param: dest_ip_address: The ip address of the machine that the payload will be sent to + :param: dest_port: The port of the machine that the payload will be sent to + :param: session_id: The id of the session + :return: True if successful, False otherwise. """ - pass + self.software_manager.send_payload_to_session_manager( + payload=payload, dest_ip_address=dest_ip_address, dest_port=self.port, session_id=session_id + ) - def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + def receive( + self, + payload: Any, + dest_ip_address: Optional[IPv4Address] = None, + dest_port: Optional[Port] = None, + session_id: Optional[str] = None, + **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: payload: The payload to send. + :param: dest_ip_address: The ip address of the machine that the payload will be sent to + :param: dest_port: The port of the machine that the payload will be sent to + :param: session_id: The id of the session + :return: True if successful, False otherwise. """ - pass + + pass def stop(self) -> None: """Stop the service.""" diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py new file mode 100644 index 00000000..fdb3426d --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py @@ -0,0 +1,66 @@ +import sys +from ipaddress import IPv4Address + +import pytest + +from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.networks import arcd_uc2_network +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.services.dns_client import DNSClient +from primaite.simulator.system.services.dns_server import DNSServer + + +@pytest.fixture(scope="function") +def dns_server() -> Node: + node = Node(hostname="dns_server") + node.software_manager.add_service(service_class=DNSServer) + node.software_manager.services["DNSServer"].start() + return node + + +@pytest.fixture(scope="function") +def dns_client() -> Node: + node = Node(hostname="dns_client") + node.software_manager.add_service(service_class=DNSClient) + node.software_manager.services["DNSClient"].start() + return node + + +def test_create_dns_server(dns_server): + assert dns_server is not None + dns_server_service: DNSServer = dns_server.software_manager.services["DNSServer"] + assert dns_server_service.name is "DNSServer" + assert dns_server_service.port is Port.DNS + assert dns_server_service.protocol is IPProtocol.UDP + + +def test_create_dns_client(dns_client): + assert dns_client is not None + dns_client_service: DNSClient = dns_client.software_manager.services["DNSClient"] + assert dns_client_service.name is "DNSClient" + assert dns_client_service.port is Port.DNS + assert dns_client_service.protocol is IPProtocol.UDP + + +def test_dns_server_domain_name_registration(dns_server): + """Test to check if the domain name registration works.""" + dns_server_service: DNSServer = dns_server.software_manager.services["DNSServer"] + + # register the web server in the domain controller + dns_server_service.dns_register(domain_name="real-domain.com", domain_ip_address=IPv4Address("192.168.1.12")) + + # return none for an unknown domain + assert dns_server_service.dns_lookup("fake-domain.com") is None + assert dns_server_service.dns_lookup("real-domain.com") is not None + + +def test_dns_client_check_domain_in_cache(dns_client): + """Test to make sure that the check_domain_in_cache returns the correct values.""" + dns_client_service: DNSClient = dns_client.software_manager.services["DNSClient"] + + # add a domain to the dns client cache + dns_client_service.add_domain_to_cache("real-domain.com", IPv4Address("192.168.1.12")) + + assert dns_client_service.check_domain_in_cache("fake-domain.com") is False + assert dns_client_service.check_domain_in_cache("real-domain.com") is True From 6f2f23e04f1ba23e2944e2ef1fe22286c447d887 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 7 Sep 2023 15:59:46 +0100 Subject: [PATCH 09/16] #1752: update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14a53d73..d9700f83 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ SessionManager. 1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP) - Red Agent Services: - Data Manipulator Bot - A red agent service which sends a payload to a target machine. (By default this payload is a SQL query that breaks a database) +- DNS Services: DNS Client and DNS Server ## [2.0.0] - 2023-07-26 From 1a81285b7621a4d3a9399170f646a79e1451c503 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 12 Sep 2023 08:46:07 +0100 Subject: [PATCH 10/16] #1752: Added send+receive functionality for DNS client and server + tests + added simulation_output to gitignore --- .gitignore | 2 + src/primaite/simulator/network/networks.py | 15 +++++ .../simulator/network/protocols/dns.py | 8 +-- .../simulator/system/services/dns_client.py | 55 +++++++++++------- .../simulator/system/services/dns_server.py | 15 ++--- .../system/test_dns_client_server.py | 24 ++++++++ .../_simulator/_system/_services/test_dns.py | 58 +++++++++++++++---- 7 files changed, 129 insertions(+), 48 deletions(-) create mode 100644 tests/integration_tests/system/test_dns_client_server.py diff --git a/.gitignore b/.gitignore index ff86b65f..66d528a8 100644 --- a/.gitignore +++ b/.gitignore @@ -144,9 +144,11 @@ cython_debug/ # IDE .idea/ docs/source/primaite-dependencies.rst +.vscode/ # outputs src/primaite/outputs/ +simulation_output/ # benchmark session outputs benchmark/output diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index ce1ef338..79af75e4 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -10,6 +10,8 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.applications.database_client import DatabaseClient from primaite.simulator.system.services.database_service import DatabaseService +from primaite.simulator.system.services.dns_client import DNSClient +from primaite.simulator.system.services.dns_server import DNSServer from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot @@ -129,6 +131,7 @@ def arcd_uc2_network() -> Network: hostname="client_1", ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" ) client_1.power_on() + client_1.software_manager.install(DNSClient) network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) client_1.software_manager.install(DataManipulationBot) db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"] @@ -139,6 +142,7 @@ def arcd_uc2_network() -> Network: hostname="client_2", ip_address="192.168.10.22", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" ) client_2.power_on() + client_2.software_manager.install(DNSClient) network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2]) # Domain Controller @@ -149,6 +153,8 @@ def arcd_uc2_network() -> Network: default_gateway="192.168.1.1", ) domain_controller.power_on() + domain_controller.software_manager.install(DNSServer) + network.connect(endpoint_b=domain_controller.ethernet_port[1], endpoint_a=switch_1.switch_ports[1]) # Database Server @@ -200,12 +206,17 @@ def arcd_uc2_network() -> Network: ) web_server.power_on() web_server.software_manager.install(DatabaseClient) + database_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] database_client.configure(server_ip_address=IPv4Address("192.168.1.14")) network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2]) database_client.run() database_client.connect() + # register the web_server to a domain + dns_server_service: DNSServer = domain_controller.software_manager.software["DNSServer"] # noqa + dns_server_service.dns_register("arcd.com", web_server.ip_address) + # Backup Server backup_server = Server( hostname="backup_server", ip_address="192.168.1.16", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" @@ -229,6 +240,10 @@ def arcd_uc2_network() -> Network: router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) + # Allow PostgreSQL requests router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER) + # Allow DNS requests + router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS) + return network diff --git a/src/primaite/simulator/network/protocols/dns.py b/src/primaite/simulator/network/protocols/dns.py index 0afa6405..e5602c91 100644 --- a/src/primaite/simulator/network/protocols/dns.py +++ b/src/primaite/simulator/network/protocols/dns.py @@ -56,7 +56,7 @@ class DNSPacket(BaseModel): :param domain_ip_address: The IP address that was being sought after from the original target domain name. :return: A new instance of DNSPacket. """ - return DNSPacket( - dns_request=DNSRequest(domain_name_request=self.dns_request.domain_name_request), - dns_reply=DNSReply(domain_name_ip_address=domain_ip_address), - ) + if domain_ip_address is not None: + self.dns_reply = DNSReply(domain_name_ip_address=IPv4Address(domain_ip_address)) + + return self diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index 3929065d..db01c05c 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -1,11 +1,15 @@ from ipaddress import IPv4Address from typing import Any, Dict, Optional +from primaite import getLogger from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.services.service import Service +_LOGGER = getLogger(__name__) + class DNSClient(Service): """Represents a DNS Client as a Service.""" @@ -52,15 +56,15 @@ class DNSClient(Service): """ self.dns_cache[domain_name] = ip_address - def check_domain_in_cache( + def check_domain_exists( self, target_domain: str, dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, + dest_port: Optional[Port] = Port.DNS, session_id: Optional[str] = None, is_reattempt: bool = False, ) -> bool: - """Function to check if domain name is in DNS client cache. + """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. @@ -80,21 +84,26 @@ class DNSClient(Service): return False else: # send a request to check if domain name exists in the DNS Server - self.send(payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id) - # call function again - return self.check_domain_in_cache( - target_domain=target_domain, + self.software_manager.send_payload_to_session_manager( + payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, - session_id=session_id, - is_reattempt=True, ) + # check if the domain has been added to cache + if self.dns_cache.get(target_domain) is None: + # call function again + return self.check_domain_exists( + target_domain=target_domain, + dest_ip_address=dest_ip_address, + dest_port=dest_port, + session_id=session_id, + is_reattempt=True, + ) + def send( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -112,15 +121,12 @@ class DNSClient(Service): :return: True if successful, False otherwise. """ # create DNS request packet - self.software_manager.send_payload_to_session_manager( - payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id - ) + software_manager: SoftwareManager = self.software_manager + software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id) def receive( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -131,11 +137,18 @@ class DNSClient(Service): is generated should be implemented in subclasses. :param payload: The payload to be sent. - :param dest_ip_address: The ip address of the payload destination. - :param dest_port: The port of the payload destination. :param session_id: The Session ID the payload is to originate from. Optional. :return: True if successful, False otherwise. """ - super().send() - # check the DNS packet (dns request, dns reply) here and see if it actually worked - pass + # The payload should be a DNS packet + if not isinstance(payload, DNSPacket): + _LOGGER.debug(f"{payload} is not a DNSPacket") + return False + # cast payload into a DNS packet + payload: DNSPacket = payload + if payload.dns_reply is not None: + # add the IP address to the client cache + self.dns_cache[payload.dns_request.domain_name_request] = payload.dns_reply.domain_name_ip_address + return True + + return False diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index 3dcd89f9..b879d515 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -75,8 +75,6 @@ class DNSServer(Service): def send( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -87,14 +85,13 @@ class DNSServer(Service): is generated should be implemented in subclasses. :param: payload: The payload to send. - :param: dest_ip_address: The ip address of the machine that the payload will be sent to - :param: dest_port: The port of the machine that the payload will be sent to :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 @@ -102,8 +99,6 @@ class DNSServer(Service): def receive( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -114,11 +109,9 @@ class DNSServer(Service): is generated should be implemented in subclasses. :param: payload: The payload to send. - :param: dest_ip_address: The ip address of the machine that the payload will be sent to - :param: dest_port: The port of the machine that the payload will be sent to :param: session_id: The id of the session - :return: True if successful, False otherwise. + :return: True if DNS request returns a valid IP, otherwise, False """ # The payload should be a DNS packet if not isinstance(payload, DNSPacket): @@ -128,10 +121,10 @@ class DNSServer(Service): payload: DNSPacket = payload if payload.dns_request is not None: # generate a reply with the correct DNS IP address - payload.generate_reply(self.dns_lookup(payload.dns_request.domain_name_request)) + payload = payload.generate_reply(self.dns_lookup(payload.dns_request.domain_name_request)) # send reply self.send(payload, session_id) - return True + return payload.dns_reply is not None return False diff --git a/tests/integration_tests/system/test_dns_client_server.py b/tests/integration_tests/system/test_dns_client_server.py new file mode 100644 index 00000000..77fa6017 --- /dev/null +++ b/tests/integration_tests/system/test_dns_client_server.py @@ -0,0 +1,24 @@ +from ipaddress import IPv4Address + +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.server import Server +from primaite.simulator.system.services.dns_client import DNSClient +from primaite.simulator.system.services.dns_server import DNSServer + + +def test_dns_client_server(uc2_network): + client_1: Computer = uc2_network.get_node_by_hostname("client_1") + domain_controller: Server = uc2_network.get_node_by_hostname("domain_controller") + + dns_client: DNSClient = client_1.software_manager.software["DNSClient"] + dns_server: DNSServer = domain_controller.software_manager.software["DNSServer"] + + # register a domain to web server + dns_server.dns_register("real-domain.com", IPv4Address("192.168.1.12")) + + dns_server.show() + + dns_client.check_domain_exists(target_domain="real-domain.com", dest_ip_address=IPv4Address("192.168.1.14")) + + # should register the domain in the client cache + assert dns_client.dns_cache.get("real-domain.com") is not None diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py index fdb3426d..943d3265 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py @@ -1,10 +1,9 @@ -import sys from ipaddress import IPv4Address import pytest from primaite.simulator.network.hardware.base import Node -from primaite.simulator.network.networks import arcd_uc2_network +from primaite.simulator.network.protocols.dns import DNSPacket, DNSReply, DNSRequest from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.services.dns_client import DNSClient @@ -14,22 +13,22 @@ from primaite.simulator.system.services.dns_server import DNSServer @pytest.fixture(scope="function") def dns_server() -> Node: node = Node(hostname="dns_server") - node.software_manager.add_service(service_class=DNSServer) - node.software_manager.services["DNSServer"].start() + node.software_manager.install(software_class=DNSServer) + node.software_manager.software["DNSServer"].start() return node @pytest.fixture(scope="function") def dns_client() -> Node: node = Node(hostname="dns_client") - node.software_manager.add_service(service_class=DNSClient) - node.software_manager.services["DNSClient"].start() + node.software_manager.install(software_class=DNSClient) + node.software_manager.software["DNSClient"].start() return node def test_create_dns_server(dns_server): assert dns_server is not None - dns_server_service: DNSServer = dns_server.software_manager.services["DNSServer"] + dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"] assert dns_server_service.name is "DNSServer" assert dns_server_service.port is Port.DNS assert dns_server_service.protocol is IPProtocol.UDP @@ -37,7 +36,7 @@ def test_create_dns_server(dns_server): def test_create_dns_client(dns_client): assert dns_client is not None - dns_client_service: DNSClient = dns_client.software_manager.services["DNSClient"] + dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"] assert dns_client_service.name is "DNSClient" assert dns_client_service.port is Port.DNS assert dns_client_service.protocol is IPProtocol.UDP @@ -45,7 +44,7 @@ def test_create_dns_client(dns_client): def test_dns_server_domain_name_registration(dns_server): """Test to check if the domain name registration works.""" - dns_server_service: DNSServer = dns_server.software_manager.services["DNSServer"] + dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"] # register the web server in the domain controller dns_server_service.dns_register(domain_name="real-domain.com", domain_ip_address=IPv4Address("192.168.1.12")) @@ -57,10 +56,45 @@ def test_dns_server_domain_name_registration(dns_server): def test_dns_client_check_domain_in_cache(dns_client): """Test to make sure that the check_domain_in_cache returns the correct values.""" - dns_client_service: DNSClient = dns_client.software_manager.services["DNSClient"] + dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"] # add a domain to the dns client cache dns_client_service.add_domain_to_cache("real-domain.com", IPv4Address("192.168.1.12")) - assert dns_client_service.check_domain_in_cache("fake-domain.com") is False - assert dns_client_service.check_domain_in_cache("real-domain.com") is True + assert dns_client_service.check_domain_exists("fake-domain.com") is False + assert dns_client_service.check_domain_exists("real-domain.com") is True + + +def test_dns_server_receive(dns_server): + """Test to make sure that the DNS Server correctly responds to a DNS Client request.""" + dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"] + + # register the web server in the domain controller + dns_server_service.dns_register(domain_name="real-domain.com", domain_ip_address=IPv4Address("192.168.1.12")) + + assert ( + dns_server_service.receive(payload=DNSPacket(dns_request=DNSRequest(domain_name_request="fake-domain.com"))) + is False + ) + + assert ( + dns_server_service.receive(payload=DNSPacket(dns_request=DNSRequest(domain_name_request="real-domain.com"))) + is True + ) + + dns_server_service.show() + + +def test_dns_client_receive(dns_client): + """Test to make sure the DNS Client knows how to deal with request responses.""" + dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"] + + dns_client_service.receive( + payload=DNSPacket( + dns_request=DNSRequest(domain_name_request="real-domain.com"), + dns_reply=DNSReply(domain_name_ip_address=IPv4Address("192.168.1.12")), + ) + ) + + # domain name should be saved to cache + assert dns_client_service.dns_cache["real-domain.com"] == IPv4Address("192.168.1.12") From fb96ef18c0dfb9a7bf0c9aac3c8e25de559cc55b Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 12 Sep 2023 09:32:28 +0100 Subject: [PATCH 11/16] #1752: remove unnecessary changes --- .../simulator/system/core/session_manager.py | 2 +- src/primaite/simulator/system/services/service.py | 14 +------------- 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index f8e97442..06701546 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -94,7 +94,7 @@ class SessionManager: @staticmethod def _get_session_key( frame: Frame, inbound_frame: bool = True - ) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]: + ) -> Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]]: """ Extracts the session key from the given frame. diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index f7f189f1..20b92027 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -1,10 +1,8 @@ from enum import Enum -from ipaddress import IPv4Address from typing import Any, Dict, Optional from primaite import getLogger from primaite.simulator.core import Action, ActionManager -from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.software import IOSoftware _LOGGER = getLogger(__name__) @@ -77,8 +75,6 @@ class Service(IOSoftware): def send( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -89,21 +85,15 @@ class Service(IOSoftware): is generated should be implemented in subclasses. :param: payload: The payload to send. - :param: dest_ip_address: The ip address of the machine that the payload will be sent to - :param: dest_port: The port of the machine that the payload will be sent to :param: session_id: The id of the session :return: True if successful, False otherwise. """ - self.software_manager.send_payload_to_session_manager( - payload=payload, dest_ip_address=dest_ip_address, dest_port=self.port, session_id=session_id - ) + self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id) def receive( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = None, session_id: Optional[str] = None, **kwargs, ) -> bool: @@ -114,8 +104,6 @@ class Service(IOSoftware): is generated should be implemented in subclasses. :param: payload: The payload to send. - :param: dest_ip_address: The ip address of the machine that the payload will be sent to - :param: dest_port: The port of the machine that the payload will be sent to :param: session_id: The id of the session :return: True if successful, False otherwise. From 8b6bc843216c148723cae97de05513b8a2713c01 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 12 Sep 2023 13:37:11 +0100 Subject: [PATCH 12/16] #1752: simplifying the DNS implementation - switch to TCP + fixing the DNS integration test --- src/primaite/simulator/network/hardware/base.py | 3 ++- src/primaite/simulator/network/networks.py | 5 +++++ .../network/transmission/data_link_layer.py | 3 --- .../simulator/system/core/session_manager.py | 6 ++++-- .../simulator/system/services/dns_client.py | 8 +++++--- .../simulator/system/services/dns_server.py | 3 ++- .../system/test_dns_client_server.py | 16 +++++++++++----- .../_simulator/_system/_services/test_dns.py | 4 ++-- 8 files changed, 31 insertions(+), 17 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index bceb385c..04262037 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -406,7 +406,8 @@ class SwitchPort(SimComponent): if self.enabled: frame.decrement_ttl() self.pcap.capture(frame) - self.connected_node.forward_frame(frame=frame, incoming_port=self) + connected_node: Node = self.connected_node + connected_node.forward_frame(frame=frame, incoming_port=self) return True return False diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 79af75e4..0b9a2299 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -132,6 +132,8 @@ def arcd_uc2_network() -> Network: ) client_1.power_on() client_1.software_manager.install(DNSClient) + client_1_dns_client_service: DNSServer = client_1.software_manager.software["DNSClient"] # noqa + client_1_dns_client_service.start() network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1]) client_1.software_manager.install(DataManipulationBot) db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"] @@ -143,6 +145,8 @@ def arcd_uc2_network() -> Network: ) client_2.power_on() client_2.software_manager.install(DNSClient) + client_2_dns_client_service: DNSServer = client_2.software_manager.software["DNSClient"] # noqa + client_2_dns_client_service.start() network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2]) # Domain Controller @@ -215,6 +219,7 @@ def arcd_uc2_network() -> Network: # register the web_server to a domain dns_server_service: DNSServer = domain_controller.software_manager.software["DNSServer"] # noqa + dns_server_service.start() dns_server_service.dns_register("arcd.com", web_server.ip_address) # Backup Server diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index 5c09210a..b7986622 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -5,7 +5,6 @@ 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 @@ -97,8 +96,6 @@ 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/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index 06701546..95ece9f9 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -74,7 +74,9 @@ class SessionManager: """ def __init__(self, sys_log: SysLog, arp_cache: "ARPCache"): - self.sessions_by_key: Dict[Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]], Session] = {} + self.sessions_by_key: Dict[ + Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]], Session + ] = {} self.sessions_by_uuid: Dict[str, Session] = {} self.sys_log: SysLog = sys_log self.software_manager: SoftwareManager = None # Noqa @@ -94,7 +96,7 @@ class SessionManager: @staticmethod def _get_session_key( frame: Frame, inbound_frame: bool = True - ) -> Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]]: + ) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]: """ Extracts the session key from the given frame. diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index db01c05c..d6e4a05b 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -22,7 +22,8 @@ class DNSClient(Service): kwargs["port"] = Port.DNS # DNS uses UDP by default # it switches to TCP when the bytes exceed 512 (or 4096) bytes - kwargs["protocol"] = IPProtocol.UDP + # TCP for now + kwargs["protocol"] = IPProtocol.TCP super().__init__(**kwargs) def describe_state(self) -> Dict: @@ -84,14 +85,15 @@ class DNSClient(Service): return False else: # send a request to check if domain name exists in the DNS Server - self.software_manager.send_payload_to_session_manager( + software_manager: SoftwareManager = self.software_manager + software_manager.send_payload_to_session_manager( payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, ) # check if the domain has been added to cache - if self.dns_cache.get(target_domain) is None: + if self.dns_cache.get(target_domain, None) is None: # call function again return self.check_domain_exists( target_domain=target_domain, diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index b879d515..c36c7034 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -23,7 +23,8 @@ class DNSServer(Service): kwargs["port"] = Port.DNS # DNS uses UDP by default # it switches to TCP when the bytes exceed 512 (or 4096) bytes - kwargs["protocol"] = IPProtocol.UDP + # TCP for now + kwargs["protocol"] = IPProtocol.TCP super().__init__(**kwargs) def describe_state(self) -> Dict: diff --git a/tests/integration_tests/system/test_dns_client_server.py b/tests/integration_tests/system/test_dns_client_server.py index 77fa6017..a4514bad 100644 --- a/tests/integration_tests/system/test_dns_client_server.py +++ b/tests/integration_tests/system/test_dns_client_server.py @@ -4,6 +4,7 @@ from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.server import Server from primaite.simulator.system.services.dns_client import DNSClient from primaite.simulator.system.services.dns_server import DNSServer +from primaite.simulator.system.services.service import ServiceOperatingState def test_dns_client_server(uc2_network): @@ -13,12 +14,17 @@ def test_dns_client_server(uc2_network): dns_client: DNSClient = client_1.software_manager.software["DNSClient"] dns_server: DNSServer = domain_controller.software_manager.software["DNSServer"] - # register a domain to web server - dns_server.dns_register("real-domain.com", IPv4Address("192.168.1.12")) + assert dns_client.operating_state == ServiceOperatingState.RUNNING + assert dns_server.operating_state == ServiceOperatingState.RUNNING dns_server.show() - dns_client.check_domain_exists(target_domain="real-domain.com", dest_ip_address=IPv4Address("192.168.1.14")) + # fake domain should not be added to dns cache + dns_client.check_domain_exists( + target_domain="fake-domain.com", dest_ip_address=IPv4Address(domain_controller.ip_address) + ) + assert dns_client.dns_cache.get("fake-domain.com", None) is None - # should register the domain in the client cache - assert dns_client.dns_cache.get("real-domain.com") is not None + # arcd.com is registered in dns server and should be saved to cache + dns_client.check_domain_exists(target_domain="arcd.com", dest_ip_address=IPv4Address(domain_controller.ip_address)) + assert dns_client.dns_cache.get("arcd.com", None) is not None diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py index 943d3265..b4f20539 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py @@ -31,7 +31,7 @@ def test_create_dns_server(dns_server): dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"] assert dns_server_service.name is "DNSServer" assert dns_server_service.port is Port.DNS - assert dns_server_service.protocol is IPProtocol.UDP + assert dns_server_service.protocol is IPProtocol.TCP def test_create_dns_client(dns_client): @@ -39,7 +39,7 @@ def test_create_dns_client(dns_client): dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"] assert dns_client_service.name is "DNSClient" assert dns_client_service.port is Port.DNS - assert dns_client_service.protocol is IPProtocol.UDP + assert dns_client_service.protocol is IPProtocol.TCP def test_dns_server_domain_name_registration(dns_server): From b0478f4e88de2a8cf8e3239b711647c593177867 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 13 Sep 2023 08:46:22 +0100 Subject: [PATCH 13/16] #1752: added positions to ACL rules for UC2 network to prevent rules being overwritten --- src/primaite/simulator/network/networks.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 0b9a2299..45056a4d 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -246,9 +246,11 @@ def arcd_uc2_network() -> Network: router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) # Allow PostgreSQL requests - router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER) + router_1.acl.add_rule( + action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0 + ) # Allow DNS requests - router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS) + router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1) return network From 98e103a984da2f5152e1e974c6981645da7cbec0 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 13 Sep 2023 09:48:38 +0100 Subject: [PATCH 14/16] #1752: added documentation for DNS Client and Server --- .../system/dns_client_server.rst | 56 +++++++++++++++++++ .../simulation_components/system/software.rst | 1 + 2 files changed, 57 insertions(+) create mode 100644 docs/source/simulation_components/system/dns_client_server.rst diff --git a/docs/source/simulation_components/system/dns_client_server.rst b/docs/source/simulation_components/system/dns_client_server.rst new file mode 100644 index 00000000..776c90b7 --- /dev/null +++ b/docs/source/simulation_components/system/dns_client_server.rst @@ -0,0 +1,56 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +DNS Client Server +====================== + +DNS Server +---------------- +Also known as a DNS Resolver, the ``DNSServer`` provides a DNS Server simulation by extending the base Service class. + +Key capabilities +^^^^^^^^^^^^^^^^ + +- Simulates DNS requests and DNSPacket transfer across a network +- Registers domain names and the IP Address linked to the domain name +- Returns the IP address for a given domain name within a DNS Packet that a DNS Client can read +- Leverages the Service base class for install/uninstall, status tracking, etc. + +Usage +^^^^^ +- Install on a Node via the ``SoftwareManager`` to start the database service. +- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future) + +Implementation +^^^^^^^^^^^^^^ + +- DNS request and responses use a ``DNSPacket`` object +- Extends Service class for integration with ``SoftwareManager``. + +DNS Client +--------------- + +The DNSClient provides a client interface for connecting to the ``DNSServer``. + +Key features +^^^^^^^^^^^^ + +- Connects to the ``DNSServer`` via the ``SoftwareManager``. +- Executes DNS lookup requests and keeps a cache of known domain name IP addresses. +- Handles connection to DNSServer and querying for domain name IP addresses. + +Usage +^^^^^ + +- Install on a Node via the ``SoftwareManager`` to start the database service. +- Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future) +- Execute domain name checks with ``check_domain_exists``, providing a ``DNSServer`` ``IPv4Address``. +- ``DNSClient`` will automatically add the IP Address of the domain into its cache + +Implementation +^^^^^^^^^^^^^^ + +- Leverages ``SoftwareManager`` for sending payloads over the network. +- Provides easy interface for Nodes to find IP addresses via domain names. +- Extends base Service class. diff --git a/docs/source/simulation_components/system/software.rst b/docs/source/simulation_components/system/software.rst index d0355d3a..275fdaf9 100644 --- a/docs/source/simulation_components/system/software.rst +++ b/docs/source/simulation_components/system/software.rst @@ -17,3 +17,4 @@ Contents database_client_server data_manipulation_bot + dns_client_server From b1e46b4f9ebf3fcb84c0b96acbf4e3e9bdb1f689 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Thu, 14 Sep 2023 20:08:06 +0100 Subject: [PATCH 15/16] #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 """ From 939de40f1e2a3f5c9fef4784e30d002d9e000352 Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 18 Sep 2023 10:25:26 +0100 Subject: [PATCH 16/16] #1752 - Moved dns_server ip address from the NIC to the Node. Updated the arcd_uc2_network so that clients and servers have a dns server. Added sys_log entries for DNSServer and DNSClient. MAde the DNSServer always rend a reply, but for the resolved IP address to be empty if it cannot be resolved. --- .../system/dns_client_server.rst | 8 ++--- .../simulator/network/hardware/base.py | 7 ++-- src/primaite/simulator/network/networks.py | 26 +++++++++++--- .../simulator/network/protocols/dns.py | 5 ++- .../simulator/system/core/software_manager.py | 13 +++++-- .../simulator/system/services/dns_client.py | 36 +++++++++---------- .../simulator/system/services/dns_server.py | 15 +++++--- .../system/test_dns_client_server.py | 6 ++-- 8 files changed, 72 insertions(+), 44 deletions(-) diff --git a/docs/source/simulation_components/system/dns_client_server.rst b/docs/source/simulation_components/system/dns_client_server.rst index 776c90b7..f57f903b 100644 --- a/docs/source/simulation_components/system/dns_client_server.rst +++ b/docs/source/simulation_components/system/dns_client_server.rst @@ -3,10 +3,10 @@ © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK DNS Client Server -====================== +================= DNS Server ----------------- +---------- Also known as a DNS Resolver, the ``DNSServer`` provides a DNS Server simulation by extending the base Service class. Key capabilities @@ -29,7 +29,7 @@ Implementation - Extends Service class for integration with ``SoftwareManager``. DNS Client ---------------- +---------- The DNSClient provides a client interface for connecting to the ``DNSServer``. @@ -45,7 +45,7 @@ Usage - Install on a Node via the ``SoftwareManager`` to start the database service. - Service runs on TCP port 53 by default. (TODO: TCP for now, should be UDP in future) -- Execute domain name checks with ``check_domain_exists``, providing a ``DNSServer`` ``IPv4Address``. +- Execute domain name checks with ``check_domain_exists``. - ``DNSClient`` will automatically add the IP Address of the domain into its cache Implementation diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 04262037..dd2130d2 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -5,7 +5,7 @@ import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Tuple, Union +from typing import Any, Dict, Literal, Optional, Tuple, Union from prettytable import MARKDOWN, PrettyTable @@ -89,8 +89,6 @@ class NIC(SimComponent): "The Maximum Transmission Unit (MTU) of the NIC in Bytes. Default is 1500 B" wake_on_lan: bool = False "Indicates if the NIC supports Wake-on-LAN functionality." - dns_servers: List[IPv4Address] = [] - "List of IP addresses of DNS servers used for name resolution." connected_node: Optional[Node] = None "The Node to which the NIC is connected." connected_link: Optional[Link] = None @@ -882,6 +880,8 @@ class Node(SimComponent): "The NICs on the node." ethernet_port: Dict[int, NIC] = {} "The NICs on the node by port id." + dns_server: Optional[IPv4Address] = None + "List of IP addresses of DNS servers used for name resolution." accounts: Dict[str, Account] = {} "All accounts on the node." @@ -931,6 +931,7 @@ class Node(SimComponent): sys_log=kwargs.get("sys_log"), session_manager=kwargs.get("session_manager"), file_system=kwargs.get("file_system"), + dns_server=kwargs.get("dns_server"), ) super().__init__(**kwargs) self.arp.nics = self.nics diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 45056a4d..78d2e68f 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -128,7 +128,11 @@ def arcd_uc2_network() -> Network: # Client 1 client_1 = Computer( - hostname="client_1", ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" + hostname="client_1", + ip_address="192.168.10.21", + subnet_mask="255.255.255.0", + default_gateway="192.168.10.1", + dns_server=IPv4Address("192.168.1.10"), ) client_1.power_on() client_1.software_manager.install(DNSClient) @@ -141,7 +145,11 @@ def arcd_uc2_network() -> Network: # Client 2 client_2 = Computer( - hostname="client_2", ip_address="192.168.10.22", subnet_mask="255.255.255.0", default_gateway="192.168.10.1" + hostname="client_2", + ip_address="192.168.10.22", + subnet_mask="255.255.255.0", + default_gateway="192.168.10.1", + dns_server=IPv4Address("192.168.1.10"), ) client_2.power_on() client_2.software_manager.install(DNSClient) @@ -167,6 +175,7 @@ def arcd_uc2_network() -> Network: ip_address="192.168.1.14", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", + dns_server=IPv4Address("192.168.1.10"), ) database_server.power_on() network.connect(endpoint_b=database_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[3]) @@ -206,7 +215,11 @@ def arcd_uc2_network() -> Network: # Web Server web_server = Server( - hostname="web_server", ip_address="192.168.1.12", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + hostname="web_server", + ip_address="192.168.1.12", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1", + dns_server=IPv4Address("192.168.1.10"), ) web_server.power_on() web_server.software_manager.install(DatabaseClient) @@ -224,7 +237,11 @@ def arcd_uc2_network() -> Network: # Backup Server backup_server = Server( - hostname="backup_server", ip_address="192.168.1.16", subnet_mask="255.255.255.0", default_gateway="192.168.1.1" + hostname="backup_server", + ip_address="192.168.1.16", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1", + dns_server=IPv4Address("192.168.1.10"), ) backup_server.power_on() network.connect(endpoint_b=backup_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[4]) @@ -235,6 +252,7 @@ def arcd_uc2_network() -> Network: ip_address="192.168.1.110", subnet_mask="255.255.255.0", default_gateway="192.168.1.1", + dns_server=IPv4Address("192.168.1.10"), ) security_suite.power_on() network.connect(endpoint_b=security_suite.ethernet_port[1], endpoint_a=switch_1.switch_ports[7]) diff --git a/src/primaite/simulator/network/protocols/dns.py b/src/primaite/simulator/network/protocols/dns.py index e5602c91..41bf5e0c 100644 --- a/src/primaite/simulator/network/protocols/dns.py +++ b/src/primaite/simulator/network/protocols/dns.py @@ -22,7 +22,7 @@ class DNSReply(BaseModel): :param domain_name_ip_address: IP Address of the Domain Name requested. """ - domain_name_ip_address: IPv4Address + domain_name_ip_address: Optional[IPv4Address] = None "IP Address of the Domain Name requested." @@ -56,7 +56,6 @@ class DNSPacket(BaseModel): :param domain_ip_address: The IP address that was being sought after from the original target domain name. :return: A new instance of DNSPacket. """ - if domain_ip_address is not None: - self.dns_reply = DNSReply(domain_name_ip_address=IPv4Address(domain_ip_address)) + self.dns_reply = DNSReply(domain_name_ip_address=domain_ip_address) return self diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index d5fda1b3..99445bf8 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -23,7 +23,13 @@ IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware) class SoftwareManager: """A class that manages all running Services and Applications on a Node and facilitates their communication.""" - def __init__(self, session_manager: "SessionManager", sys_log: SysLog, file_system: FileSystem): + def __init__( + self, + session_manager: "SessionManager", + sys_log: SysLog, + file_system: FileSystem, + dns_server: Optional[IPv4Address], + ): """ Initialize a new instance of SoftwareManager. @@ -35,6 +41,7 @@ class SoftwareManager: self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {} self.sys_log: SysLog = sys_log self.file_system: FileSystem = file_system + self.dns_server: Optional[IPv4Address] = dns_server def get_open_ports(self) -> List[Port]: """ @@ -58,7 +65,9 @@ class SoftwareManager: if software_class in self._software_class_to_name_map: self.sys_log.info(f"Cannot install {software_class} as it is already installed") return - software = software_class(software_manager=self, sys_log=self.sys_log, file_system=self.file_system) + software = software_class( + software_manager=self, sys_log=self.sys_log, file_system=self.file_system, dns_server=self.dns_server + ) if isinstance(software, Application): software.install() software.software_manager = self diff --git a/src/primaite/simulator/system/services/dns_client.py b/src/primaite/simulator/system/services/dns_client.py index af20d6b8..cf5278af 100644 --- a/src/primaite/simulator/system/services/dns_client.py +++ b/src/primaite/simulator/system/services/dns_client.py @@ -16,6 +16,8 @@ class DNSClient(Service): dns_cache: Dict[str, IPv4Address] = {} "A dict of known mappings between domain/URLs names and IPv4 addresses." + dns_server: Optional[IPv4Address] = None + "The DNS Server the client sends requests to." def __init__(self, **kwargs): kwargs["name"] = "DNSClient" @@ -60,16 +62,12 @@ class DNSClient(Service): def check_domain_exists( self, target_domain: str, - dest_ip_address: Optional[IPv4Address] = None, - dest_port: Optional[Port] = Port.DNS, session_id: Optional[str] = None, is_reattempt: bool = False, ) -> bool: """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 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. """ @@ -78,30 +76,28 @@ class DNSClient(Service): # check if the domain is already in the DNS cache if target_domain in self.dns_cache: + self.sys_log.info( + f"DNS Client: Domain lookup for {target_domain} successful, resolves to {self.dns_cache[target_domain]}" + ) return True else: # return False if already reattempted if is_reattempt: + self.sys_log.info(f"DNS Client: Domain lookup for {target_domain} failed") return False else: # send a request to check if domain name exists in the DNS Server software_manager: SoftwareManager = self.software_manager software_manager.send_payload_to_session_manager( - payload=payload, - dest_ip_address=dest_ip_address, - dest_port=dest_port, + payload=payload, dest_ip_address=self.dns_server, dest_port=Port.DNS ) - # check if the domain has been added to cache - if self.dns_cache.get(target_domain, None) is None: - # call function again - return self.check_domain_exists( - target_domain=target_domain, - dest_ip_address=dest_ip_address, - dest_port=dest_port, - session_id=session_id, - is_reattempt=True, - ) + # recursively re-call the function passing is_reattempt=True + return self.check_domain_exists( + target_domain=target_domain, + session_id=session_id, + is_reattempt=True, + ) def send( self, @@ -125,6 +121,7 @@ class DNSClient(Service): # create DNS request packet software_manager: SoftwareManager = self.software_manager software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id) + return True def receive( self, @@ -150,7 +147,8 @@ class DNSClient(Service): payload: DNSPacket = payload if payload.dns_reply is not None: # add the IP address to the client cache - self.dns_cache[payload.dns_request.domain_name_request] = payload.dns_reply.domain_name_ip_address - return True + if payload.dns_reply.domain_name_ip_address: + self.dns_cache[payload.dns_request.domain_name_request] = payload.dns_reply.domain_name_ip_address + return True return False diff --git a/src/primaite/simulator/system/services/dns_server.py b/src/primaite/simulator/system/services/dns_server.py index be1145a2..c6a9afd3 100644 --- a/src/primaite/simulator/system/services/dns_server.py +++ b/src/primaite/simulator/system/services/dns_server.py @@ -47,10 +47,7 @@ class DNSServer(Service): :param target_domain: The single domain name requested by a DNS client. :return ip_address: The IP address of that domain name or None. """ - if target_domain in self.dns_table: - return self.dns_table[target_domain] - else: - return None + return self.dns_table.get(target_domain) def dns_register(self, domain_name: str, domain_ip_address: IPv4Address): """ @@ -97,11 +94,19 @@ class DNSServer(Service): # cast payload into a DNS packet payload: DNSPacket = payload if payload.dns_request is not None: + self.sys_log.info( + f"DNS Server: Received domain lookup request for {payload.dns_request.domain_name_request} " + f"from session {session_id}" + ) # generate a reply with the correct DNS IP address payload = payload.generate_reply(self.dns_lookup(payload.dns_request.domain_name_request)) + self.sys_log.info( + f"DNS Server: Responding to domain lookup request for {payload.dns_request.domain_name_request} " + f"with ip address: {payload.dns_reply.domain_name_ip_address}" + ) # send reply self.send(payload, session_id) - return payload.dns_reply is not None + return payload.dns_reply.domain_name_ip_address is not None return False diff --git a/tests/integration_tests/system/test_dns_client_server.py b/tests/integration_tests/system/test_dns_client_server.py index a4514bad..640c268a 100644 --- a/tests/integration_tests/system/test_dns_client_server.py +++ b/tests/integration_tests/system/test_dns_client_server.py @@ -20,11 +20,9 @@ def test_dns_client_server(uc2_network): dns_server.show() # fake domain should not be added to dns cache - dns_client.check_domain_exists( - target_domain="fake-domain.com", dest_ip_address=IPv4Address(domain_controller.ip_address) - ) + assert not dns_client.check_domain_exists(target_domain="fake-domain.com") assert dns_client.dns_cache.get("fake-domain.com", None) is None # arcd.com is registered in dns server and should be saved to cache - dns_client.check_domain_exists(target_domain="arcd.com", dest_ip_address=IPv4Address(domain_controller.ip_address)) + assert dns_client.check_domain_exists(target_domain="arcd.com") assert dns_client.dns_cache.get("arcd.com", None) is not None