codekingpro/portable-devtools

1import sys2 3import dns._features4 5if sys.platform == "win32":6    from typing import Any7 8    import dns.name9 10    _prefer_wmi = True11 12    import winreg  # pylint: disable=import-error13 14    # Keep pylint quiet on non-windows.15    try:16        WindowsError is None  # pylint: disable=used-before-assignment17    except KeyError:18        WindowsError = Exception19 20    if dns._features.have("wmi"):21        import threading22 23        import pythoncom  # pylint: disable=import-error24        import wmi  # pylint: disable=import-error25 26        _have_wmi = True27    else:28        _have_wmi = False29 30    def _config_domain(domain):31        # Sometimes DHCP servers add a '.' prefix to the default domain, and32        # Windows just stores such values in the registry (see #687).33        # Check for this and fix it.34        if domain.startswith("."):35            domain = domain[1:]36        return dns.name.from_text(domain)37 38    class DnsInfo:39        def __init__(self):40            self.domain = None41            self.nameservers = []42            self.search = []43 44    if _have_wmi:45 46        class _WMIGetter(threading.Thread):47            def __init__(self):48                super().__init__()49                self.info = DnsInfo()50 51            def run(self):52                pythoncom.CoInitialize()53                try:54                    system = wmi.WMI()55                    for interface in system.Win32_NetworkAdapterConfiguration():56                        if interface.IPEnabled and interface.DNSServerSearchOrder:57                            self.info.nameservers = list(interface.DNSServerSearchOrder)58                            if interface.DNSDomain:59                                self.info.domain = _config_domain(interface.DNSDomain)60                            if interface.DNSDomainSuffixSearchOrder:61                                self.info.search = [62                                    _config_domain(x)63                                    for x in interface.DNSDomainSuffixSearchOrder64                                ]65                            break66                finally:67                    pythoncom.CoUninitialize()68 69            def get(self):70                # We always run in a separate thread to avoid any issues with71                # the COM threading model.72                self.start()73                self.join()74                return self.info75 76    else:77 78        class _WMIGetter:  # type: ignore79            pass80 81    class _RegistryGetter:82        def __init__(self):83            self.info = DnsInfo()84 85        def _determine_split_char(self, entry):86            #87            # The windows registry irritatingly changes the list element88            # delimiter in between ' ' and ',' (and vice-versa) in various89            # versions of windows.90            #91            if entry.find(" ") >= 0:92                split_char = " "93            elif entry.find(",") >= 0:94                split_char = ","95            else:96                # probably a singleton; treat as a space-separated list.97                split_char = " "98            return split_char99 100        def _config_nameservers(self, nameservers):101            split_char = self._determine_split_char(nameservers)102            ns_list = nameservers.split(split_char)103            for ns in ns_list:104                if ns not in self.info.nameservers:105                    self.info.nameservers.append(ns)106 107        def _config_search(self, search):108            split_char = self._determine_split_char(search)109            search_list = search.split(split_char)110            for s in search_list:111                s = _config_domain(s)112                if s not in self.info.search:113                    self.info.search.append(s)114 115        def _config_fromkey(self, key, always_try_domain):116            try:117                servers, _ = winreg.QueryValueEx(key, "NameServer")118            except WindowsError:119                servers = None120            if servers:121                self._config_nameservers(servers)122            if servers or always_try_domain:123                try:124                    dom, _ = winreg.QueryValueEx(key, "Domain")125                    if dom:126                        self.info.domain = _config_domain(dom)127                except WindowsError:128                    pass129            else:130                try:131                    servers, _ = winreg.QueryValueEx(key, "DhcpNameServer")132                except WindowsError:133                    servers = None134                if servers:135                    self._config_nameservers(servers)136                    try:137                        dom, _ = winreg.QueryValueEx(key, "DhcpDomain")138                        if dom:139                            self.info.domain = _config_domain(dom)140                    except WindowsError:141                        pass142            try:143                search, _ = winreg.QueryValueEx(key, "SearchList")144            except WindowsError:145                search = None146            if search is None:147                try:148                    search, _ = winreg.QueryValueEx(key, "DhcpSearchList")149                except WindowsError:150                    search = None151            if search:152                self._config_search(search)153 154        def _is_nic_enabled(self, lm, guid):155            # Look in the Windows Registry to determine whether the network156            # interface corresponding to the given guid is enabled.157            #158            # (Code contributed by Paul Marks, thanks!)159            #160            try:161                # This hard-coded location seems to be consistent, at least162                # from Windows 2000 through Vista.163                connection_key = winreg.OpenKey(164                    lm,165                    r"SYSTEM\CurrentControlSet\Control\Network"166                    r"\{4D36E972-E325-11CE-BFC1-08002BE10318}"167                    r"\%s\Connection" % guid,168                )169 170                try:171                    # The PnpInstanceID points to a key inside Enum172                    (pnp_id, ttype) = winreg.QueryValueEx(173                        connection_key, "PnpInstanceID"174                    )175 176                    if ttype != winreg.REG_SZ:177                        raise ValueError  # pragma: no cover178 179                    device_key = winreg.OpenKey(180                        lm, r"SYSTEM\CurrentControlSet\Enum\%s" % pnp_id181                    )182 183                    try:184                        # Get ConfigFlags for this device185                        (flags, ttype) = winreg.QueryValueEx(device_key, "ConfigFlags")186 187                        if ttype != winreg.REG_DWORD:188                            raise ValueError  # pragma: no cover189 190                        # Based on experimentation, bit 0x1 indicates that the191                        # device is disabled.192                        #193                        # XXXRTH I suspect we really want to & with 0x03 so194                        # that CONFIGFLAGS_REMOVED devices are also ignored,195                        # but we're shifting to WMI as ConfigFlags is not196                        # supposed to be used.197                        return not flags & 0x1198 199                    finally:200                        device_key.Close()201                finally:202                    connection_key.Close()203            except Exception:  # pragma: no cover204                return False205 206        def get(self):207            """Extract resolver configuration from the Windows registry."""208 209            lm = winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE)210            try:211                tcp_params = winreg.OpenKey(212                    lm, r"SYSTEM\CurrentControlSet\Services\Tcpip\Parameters"213                )214                try:215                    self._config_fromkey(tcp_params, True)216                finally:217                    tcp_params.Close()218                interfaces = winreg.OpenKey(219                    lm,220                    r"SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces",221                )222                try:223                    i = 0224                    while True:225                        try:226                            guid = winreg.EnumKey(interfaces, i)227                            i += 1228                            key = winreg.OpenKey(interfaces, guid)229                            try:230                                if not self._is_nic_enabled(lm, guid):231                                    continue232                                self._config_fromkey(key, False)233                            finally:234                                key.Close()235                        except EnvironmentError:236                            break237                finally:238                    interfaces.Close()239            finally:240                lm.Close()241            return self.info242 243    _getter_class: Any244    if _have_wmi and _prefer_wmi:245        _getter_class = _WMIGetter246    else:247        _getter_class = _RegistryGetter248 249    def get_dns_info():250        """Extract resolver configuration."""251        getter = _getter_class()252        return getter.get()253