codekingpro/portable-devtools

1from typing import Optional, Union2from urllib.parse import urlparse3 4import dns.asyncbackend5import dns.asyncquery6import dns.inet7import dns.message8import dns.query9 10 11class Nameserver:12    def __init__(self):13        pass14 15    def __str__(self):16        raise NotImplementedError17 18    def kind(self) -> str:19        raise NotImplementedError20 21    def is_always_max_size(self) -> bool:22        raise NotImplementedError23 24    def answer_nameserver(self) -> str:25        raise NotImplementedError26 27    def answer_port(self) -> int:28        raise NotImplementedError29 30    def query(31        self,32        request: dns.message.QueryMessage,33        timeout: float,34        source: Optional[str],35        source_port: int,36        max_size: bool,37        one_rr_per_rrset: bool = False,38        ignore_trailing: bool = False,39    ) -> dns.message.Message:40        raise NotImplementedError41 42    async def async_query(43        self,44        request: dns.message.QueryMessage,45        timeout: float,46        source: Optional[str],47        source_port: int,48        max_size: bool,49        backend: dns.asyncbackend.Backend,50        one_rr_per_rrset: bool = False,51        ignore_trailing: bool = False,52    ) -> dns.message.Message:53        raise NotImplementedError54 55 56class AddressAndPortNameserver(Nameserver):57    def __init__(self, address: str, port: int):58        super().__init__()59        self.address = address60        self.port = port61 62    def kind(self) -> str:63        raise NotImplementedError64 65    def is_always_max_size(self) -> bool:66        return False67 68    def __str__(self):69        ns_kind = self.kind()70        return f"{ns_kind}:{self.address}@{self.port}"71 72    def answer_nameserver(self) -> str:73        return self.address74 75    def answer_port(self) -> int:76        return self.port77 78 79class Do53Nameserver(AddressAndPortNameserver):80    def __init__(self, address: str, port: int = 53):81        super().__init__(address, port)82 83    def kind(self):84        return "Do53"85 86    def query(87        self,88        request: dns.message.QueryMessage,89        timeout: float,90        source: Optional[str],91        source_port: int,92        max_size: bool,93        one_rr_per_rrset: bool = False,94        ignore_trailing: bool = False,95    ) -> dns.message.Message:96        if max_size:97            response = dns.query.tcp(98                request,99                self.address,100                timeout=timeout,101                port=self.port,102                source=source,103                source_port=source_port,104                one_rr_per_rrset=one_rr_per_rrset,105                ignore_trailing=ignore_trailing,106            )107        else:108            response = dns.query.udp(109                request,110                self.address,111                timeout=timeout,112                port=self.port,113                source=source,114                source_port=source_port,115                raise_on_truncation=True,116                one_rr_per_rrset=one_rr_per_rrset,117                ignore_trailing=ignore_trailing,118                ignore_errors=True,119                ignore_unexpected=True,120            )121        return response122 123    async def async_query(124        self,125        request: dns.message.QueryMessage,126        timeout: float,127        source: Optional[str],128        source_port: int,129        max_size: bool,130        backend: dns.asyncbackend.Backend,131        one_rr_per_rrset: bool = False,132        ignore_trailing: bool = False,133    ) -> dns.message.Message:134        if max_size:135            response = await dns.asyncquery.tcp(136                request,137                self.address,138                timeout=timeout,139                port=self.port,140                source=source,141                source_port=source_port,142                backend=backend,143                one_rr_per_rrset=one_rr_per_rrset,144                ignore_trailing=ignore_trailing,145            )146        else:147            response = await dns.asyncquery.udp(148                request,149                self.address,150                timeout=timeout,151                port=self.port,152                source=source,153                source_port=source_port,154                raise_on_truncation=True,155                backend=backend,156                one_rr_per_rrset=one_rr_per_rrset,157                ignore_trailing=ignore_trailing,158                ignore_errors=True,159                ignore_unexpected=True,160            )161        return response162 163 164class DoHNameserver(Nameserver):165    def __init__(166        self,167        url: str,168        bootstrap_address: Optional[str] = None,169        verify: Union[bool, str] = True,170        want_get: bool = False,171    ):172        super().__init__()173        self.url = url174        self.bootstrap_address = bootstrap_address175        self.verify = verify176        self.want_get = want_get177 178    def kind(self):179        return "DoH"180 181    def is_always_max_size(self) -> bool:182        return True183 184    def __str__(self):185        return self.url186 187    def answer_nameserver(self) -> str:188        return self.url189 190    def answer_port(self) -> int:191        port = urlparse(self.url).port192        if port is None:193            port = 443194        return port195 196    def query(197        self,198        request: dns.message.QueryMessage,199        timeout: float,200        source: Optional[str],201        source_port: int,202        max_size: bool = False,203        one_rr_per_rrset: bool = False,204        ignore_trailing: bool = False,205    ) -> dns.message.Message:206        return dns.query.https(207            request,208            self.url,209            timeout=timeout,210            source=source,211            source_port=source_port,212            bootstrap_address=self.bootstrap_address,213            one_rr_per_rrset=one_rr_per_rrset,214            ignore_trailing=ignore_trailing,215            verify=self.verify,216            post=(not self.want_get),217        )218 219    async def async_query(220        self,221        request: dns.message.QueryMessage,222        timeout: float,223        source: Optional[str],224        source_port: int,225        max_size: bool,226        backend: dns.asyncbackend.Backend,227        one_rr_per_rrset: bool = False,228        ignore_trailing: bool = False,229    ) -> dns.message.Message:230        return await dns.asyncquery.https(231            request,232            self.url,233            timeout=timeout,234            source=source,235            source_port=source_port,236            bootstrap_address=self.bootstrap_address,237            one_rr_per_rrset=one_rr_per_rrset,238            ignore_trailing=ignore_trailing,239            verify=self.verify,240            post=(not self.want_get),241        )242 243 244class DoTNameserver(AddressAndPortNameserver):245    def __init__(246        self,247        address: str,248        port: int = 853,249        hostname: Optional[str] = None,250        verify: Union[bool, str] = True,251    ):252        super().__init__(address, port)253        self.hostname = hostname254        self.verify = verify255 256    def kind(self):257        return "DoT"258 259    def query(260        self,261        request: dns.message.QueryMessage,262        timeout: float,263        source: Optional[str],264        source_port: int,265        max_size: bool = False,266        one_rr_per_rrset: bool = False,267        ignore_trailing: bool = False,268    ) -> dns.message.Message:269        return dns.query.tls(270            request,271            self.address,272            port=self.port,273            timeout=timeout,274            one_rr_per_rrset=one_rr_per_rrset,275            ignore_trailing=ignore_trailing,276            server_hostname=self.hostname,277            verify=self.verify,278        )279 280    async def async_query(281        self,282        request: dns.message.QueryMessage,283        timeout: float,284        source: Optional[str],285        source_port: int,286        max_size: bool,287        backend: dns.asyncbackend.Backend,288        one_rr_per_rrset: bool = False,289        ignore_trailing: bool = False,290    ) -> dns.message.Message:291        return await dns.asyncquery.tls(292            request,293            self.address,294            port=self.port,295            timeout=timeout,296            one_rr_per_rrset=one_rr_per_rrset,297            ignore_trailing=ignore_trailing,298            server_hostname=self.hostname,299            verify=self.verify,300        )301 302 303class DoQNameserver(AddressAndPortNameserver):304    def __init__(305        self,306        address: str,307        port: int = 853,308        verify: Union[bool, str] = True,309        server_hostname: Optional[str] = None,310    ):311        super().__init__(address, port)312        self.verify = verify313        self.server_hostname = server_hostname314 315    def kind(self):316        return "DoQ"317 318    def query(319        self,320        request: dns.message.QueryMessage,321        timeout: float,322        source: Optional[str],323        source_port: int,324        max_size: bool = False,325        one_rr_per_rrset: bool = False,326        ignore_trailing: bool = False,327    ) -> dns.message.Message:328        return dns.query.quic(329            request,330            self.address,331            port=self.port,332            timeout=timeout,333            one_rr_per_rrset=one_rr_per_rrset,334            ignore_trailing=ignore_trailing,335            verify=self.verify,336            server_hostname=self.server_hostname,337        )338 339    async def async_query(340        self,341        request: dns.message.QueryMessage,342        timeout: float,343        source: Optional[str],344        source_port: int,345        max_size: bool,346        backend: dns.asyncbackend.Backend,347        one_rr_per_rrset: bool = False,348        ignore_trailing: bool = False,349    ) -> dns.message.Message:350        return await dns.asyncquery.quic(351            request,352            self.address,353            port=self.port,354            timeout=timeout,355            one_rr_per_rrset=one_rr_per_rrset,356            ignore_trailing=ignore_trailing,357            verify=self.verify,358            server_hostname=self.server_hostname,359        )360