108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
# 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
|