adding dns server and zone utilities to the ad role

Author: Dan Lavu

sssd_test_framework/roles/ad.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
+from abc import ABC
 from datetime import datetime
 from typing import Any, TypeAlias
 
 from pytest_mh.cli import CLIBuilderArgs
 from pytest_mh.conn import ProcessResult
 
 from ..hosts.ad import ADHost
-from ..misc import attrs_include_value, attrs_parse, attrs_to_hash, seconds_to_timespan
+from ..misc import attrs_include_value, attrs_parse, attrs_to_hash, ip_version, seconds_to_timespan
 from .base import BaseObject, BaseWindowsRole, DeleteAttribute
 from .generic import GenericPasswordPolicy
 from .ldap import LDAPNetgroupMember
@@ -27,6 +28,8 @@
     "ADUser",
     "ADPasswordPolicy",
     "GPO",
+    "ADDNSServer",
+    "ADDNSZone",
 ]
 
 
@@ -360,6 +363,38 @@ def test_example(client: Client, ad: AD):
         """
         return ADComputer(self, name, basedn)
 
+    def dns(self) -> ADDNSServer:
+        """
+        Get DNS server object.
+
+            Get methods use dig and is parsed by jc. The data from jc contains several nested dict,
+            but two are returned as a tuple, ``answer, authority``.
+
+        .. code-block:: python
+            :caption: Example usage
+
+            # Create forward zone and add forward record
+            zone = ad.dns().zone("example.test").create()
+            zone.add_record("client", "172.16.200.15")
+
+            # Create reverse zone and add reverse record
+            zone = ad.dns().zone("10.0.10.in-addr.arpa").create()
+            zone.add_ptr_record("client.example.test", 15)
+
+            # Add forward record to default domain
+            ad.dns().zone(ad.domain).add_record("client", "1.2.3.4")
+
+            # Add a global forwarder
+            ad.dns().add_forwarder("1.1.1.1")
+
+            # Remove a global forwarder
+            ad.dns().remove_forwarder("1.1.1.1")
+
+            # Clear all forwarders
+            ad.dns().clear_forwarders()
+        """
+        return ADDNSServer(self)
+
     def gpo(self, name: str) -> GPO:
         """
         Get group policy object.
@@ -2114,4 +2149,185 @@ def lockout(self, duration: int, attempts: int) -> ADPasswordPolicy:
         return self
 
 
+class ADDNSServer(BaseObject[ADHost, AD]):
+    def __init__(self, role: AD):
+        """
+        :param role: AD host object.
+        :type role: ADHost
+        """
+        super().__init__(role)
+
+        self.domain: str = role.domain
+        """Domain name."""
+
+        self.server: str = role.server
+        """Server name."""
+
+    def zone(self, name: str) -> ADDNSZone:
+        """
+        Get ADDNsServerZone object.
+
+        :param name: Zone name.
+        :type name: str
+        :return: ADDNSServerZone object.
+        :rtype: ADDNSZone
+        """
+        return ADDNSZone(self.role, name)
+
+    def get_forwarders(self) -> list[str] | None:
+        """
+        Get DNS global forwarders.
+
+        :return: DNS global forwarders.
+        :rtype: list[str]
+        """
+        forwarders = self.host.conn.run("Get-DnsServerForwarder").stdout_lines
+        if forwarders is not None:
+            parsed_forwarders = attrs_parse(forwarders, ["IPAddress"])
+            if isinstance(parsed_forwarders, dict):
+                ip_addresses = parsed_forwarders.get("IPAddress")
+                if isinstance(ip_addresses, list) and len(ip_addresses) > 0:
+                    ip_addresses = [s.strip() for s in ip_addresses]
+                    ip_addresses = [r.replace("...", "") for r in ip_addresses]
+                    return ip_addresses[0].strip("{}").split(",")
+                else:
+                    return None
+            else:
+                return None
+
+    def add_forwarder(self, ip_address: str) -> ADDNSServer:
+        """
+        Add a DNS server forwarder.
+
+        :param ip_address: IP address.,
+        :type ip_address: str
+        :return:  Self.
+        :rtype: ADDNSServer
+        """
+        self.host.conn.run(f"Add-DnsServerForwarder -IPAddress {ip_address}")
+
+        return self
+
+    def remove_forwarder(self, ip_address: str) -> None:
+        """
+        Remove a DNS server forwarder.
+
+        :param ip_address: IP address.
+        :type ip_address: str
+        """
+        if ip_address is not None:
+            self.host.conn.run(f"Remove-DnsServerForwarder {ip_address} -Force")
+
+    def clear_forwarders(self) -> None:
+        """
+        Clear all DNS server forwarders.
+
+        AD has about four global forwarders enabled by default.
+        """
+        forwarders = self.get_forwarders()
+
+        if isinstance(forwarders, list) and not None:
+            for forwarder in forwarders:
+                self.remove_forwarder(forwarder)
+
+    def list_zones(self) -> list[str]:
+        """
+        List zones.
+        :return: List of zones.
+        :rtype: list[str]
+        """
+        result = self.host.conn.run("Get-DnsServerZone | Format-List -Property ZoneName").stdout_lines
+        result = [x for x in result if x not in ["\r", "", None]]
+        result = [y.replace("\r", "").strip() for y in result]
+        result = [z.split(":")[1].strip() for z in result]
+
+        return result
+
+
+class ADDNSZone(ADDNSServer, ABC):
+    """
+    DNS zone management.
+    """
+
+    def __init__(self, role: AD, name: str):
+        """
+        :param name: DNS zone name.
+        :type name: str
+        :param name: DNS zone name.
+        :type name: str
+        """
+        super().__init__(role)
+
+        self.zone_name: str = name
+        """Zone name."""
+
+    def create(self) -> ADDNSZone:
+        """
+        Create new zone.
+
+        :return: ADDNSServer object.
+        :rtype: ADDNSServer
+        """
+        self.host.conn.run(f"Add-DnsServerPrimaryZone -Name {self.zone_name} -ReplicationScope Forest -Passthru")
+        return self
+
+    def delete(self) -> None:
+        """
+        Delete zone.
+        """
+        self.host.conn.run(f"Delete-DnsServerZone -Name {self.zone_name}")
+
+    def add_record(self, name: str, data: str | int) -> ADDNSZone:
+        """
+        Add DNS record.
+
+        If ``data`` is a str, a forward record will be added.
+        If an integer a reverse record will be added.
+
+        :param name: Record name.
+        :type name: str
+        :param data: Record data.
+        :type data: str
+        :return: ADDNSZone object.
+        :rtype: ADDNSZone
+        """
+        args = ""
+
+        if isinstance(data, int):
+            args = f"-Ptr -Name {str(data)} -AllowUpdateAny -PtrDomainName {name}.{self.zone_name}"
+        elif isinstance(data, str) and ip_version(data) == 4:
+            args = f"-A -Name {name} -IPv4Address {data}"
+        elif isinstance(data, str) and ip_version(data) == 6:
+            args = f"-A -Name {name} -IPv6Address {data}"
+
+        self.host.conn.run(f"Add-DnsServerResourceRecord -ZoneName {self.zone_name} {args} ")
+        return self
+
+    def delete_record(self, name: str) -> None:
+        """
+        Delete DNS record.
+
+        :param name: Name of the record.
+        :type name: str
+        """
+        if "in-addr" in self.zone_name:
+            record_type = "PTR"
+        else:
+            data = self.host.conn.run(f"dig +short {name}").stdout.strip()
+            record_type = "AAAA" if ":" in data else "A"
+
+        self.host.conn.run(
+            f"Remove-DnsServerResourceRecord -ZoneName {self.zone_name} -Name {name} -RRType {record_type} -Force"
+        )
+
+    def print(self) -> str:
+        """
+        Prints all dns records in zone as text.
+
+        :return: Print zone data.
+        :rtype: str
+        """
+        return self.host.conn.run(f"Get-DnsServerResourceRecord -ZoneName {self.zone_name}").stdout
+
+
 ADNetgroupMember: TypeAlias = LDAPNetgroupMember[ADUser, ADNetgroup]