Initial commit

venv/lib/python3.8/site-packages/pymongo/srv_resolver.py

@@ -0,0 +1,107 @@
+# Copyright 2019-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you
+# may not use this file except in compliance with the License.  You
+# may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.  See the License for the specific language governing
+# permissions and limitations under the License.
+
+"""Support for resolving hosts and options from mongodb+srv:// URIs."""
+
+try:
+    from dns import resolver
+    _HAVE_DNSPYTHON = True
+except ImportError:
+    _HAVE_DNSPYTHON = False
+
+from bson.py3compat import PY3
+
+from pymongo.common import CONNECT_TIMEOUT
+from pymongo.errors import ConfigurationError
+
+
+if PY3:
+    # dnspython can return bytes or str from various parts
+    # of its API depending on version. We always want str.
+    def maybe_decode(text):
+        if isinstance(text, bytes):
+            return text.decode()
+        return text
+else:
+    def maybe_decode(text):
+        return text
+
+
+class _SrvResolver(object):
+    def __init__(self, fqdn, connect_timeout=None):
+        self.__fqdn = fqdn
+        self.__connect_timeout = connect_timeout or CONNECT_TIMEOUT
+
+        # Validate the fully qualified domain name.
+        try:
+            self.__plist = self.__fqdn.split(".")[1:]
+        except Exception:
+            raise ConfigurationError("Invalid URI host: %s" % (fqdn,))
+        self.__slen = len(self.__plist)
+        if self.__slen < 2:
+            raise ConfigurationError("Invalid URI host: %s" % (fqdn,))
+
+    def get_options(self):
+        try:
+            results = resolver.query(self.__fqdn, 'TXT',
+                                     lifetime=self.__connect_timeout)
+        except (resolver.NoAnswer, resolver.NXDOMAIN):
+            # No TXT records
+            return None
+        except Exception as exc:
+            raise ConfigurationError(str(exc))
+        if len(results) > 1:
+            raise ConfigurationError('Only one TXT record is supported')
+        return (
+            b'&'.join([b''.join(res.strings) for res in results])).decode(
+            'utf-8')
+
+    def _resolve_uri(self, encapsulate_errors):
+        try:
+            results = resolver.query('_mongodb._tcp.' + self.__fqdn, 'SRV',
+                                     lifetime=self.__connect_timeout)
+        except Exception as exc:
+            if not encapsulate_errors:
+                # Raise the original error.
+                raise
+            # Else, raise all errors as ConfigurationError.
+            raise ConfigurationError(str(exc))
+        return results
+
+    def _get_srv_response_and_hosts(self, encapsulate_errors):
+        results = self._resolve_uri(encapsulate_errors)
+
+        # Construct address tuples
+        nodes = [
+            (maybe_decode(res.target.to_text(omit_final_dot=True)), res.port)
+            for res in results]
+
+        # Validate hosts
+        for node in nodes:
+            try:
+                nlist = node[0].split(".")[1:][-self.__slen:]
+            except Exception:
+                raise ConfigurationError("Invalid SRV host: %s" % (node[0],))
+            if self.__plist != nlist:
+                raise ConfigurationError("Invalid SRV host: %s" % (node[0],))
+
+        return results, nodes
+
+    def get_hosts(self):
+        _, nodes = self._get_srv_response_and_hosts(True)
+        return nodes
+
+    def get_hosts_and_min_ttl(self):
+        results, nodes = self._get_srv_response_and_hosts(False)
+        return nodes, results.rrset.ttl