text: '''
  Test hostdb
  '''
  
  #  Licensed to the Apache Software Foundation (ASF) under one
  #  or more contributor license agreements.  See the NOTICE file
  #  distributed with this work for additional information
  #  regarding copyright ownership.  The ASF licenses this file
  #  to you 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.
  
  import os
  import requests
  import time
  import logging
  import socket
  import SocketServer
  
  import contextlib
  import dnslib
  import dnslib.server
  
  import tsqa.test_cases
  import helpers
  
  log = logging.getLogger(__name__)
  
  
  @contextlib.contextmanager
  def kill_dns(dns_server):
      ''' Temporarily kill the dns server
      '''
      dns_server.stop()
      yield
      dns_server.start_thread()
  
  
  class StubDNSResolver(object):
      '''Resolver to serve defined responses from `response_dict` or return SOA
      '''
  
      def __init__(self, responses):
          self.responses = responses
          self.resp_headers = {}
  
      def resolve(self, request, handler):
          reply = request.reply()
          for q in request.questions:
              qname = str(q.get_qname())
              if qname in self.responses:
                  for resp in self.responses[qname]:
                      reply.add_answer(resp)
              else:
                  reply.add_answer(dnslib.server.RR(
                      q.get_qname(),
                      rtype=dnslib.QTYPE.SOA,
                      ttl=1,
                      rdata=dnslib.dns.SOA(
                          'nameserver.local',
                          q.get_qname(),
                      ),
                  ))
          for k, v in self.resp_headers.iteritems():
              if k == 'rcode':
                  reply.header.set_rcode(v)
                  print 'setting rcode'
              else:
                  log.warning('Unsupported header sent to StubDNSResolver %s' % k)
          return reply
  
  
  class EchoServerIpHandler(SocketServer.BaseRequestHandler):
      """
      A subclass of RequestHandler which will return a connection uuid
      """
  
      def handle(self):
          # Receive the data in small chunks and retransmit it
          while True:
              data = self.request.recv(4096).strip()
              if data:
                  log.debug('Sending data back to the client')
              else:
                  log.debug('Client disconnected')
                  break
              resp = ('HTTP/1.1 200 OK\r\n'
                      'Content-Length: 0\r\n'
                      'Content-Type: text/html; charset=UTF-8\r\n'
                      'Connection: keep-alive\r\n'
                      'X-Server-Ip: {server_ip}\r\n'
                      'X-Server-Port: {server_port}\r\n'
                      'X-Client-Ip: {client_ip}\r\n'
                      'X-Client-Port: {client_port}\r\n'
                      '\r\n'.format(
                          server_ip=self.request.getsockname()[0],
                          server_port=self.request.getsockname()[1],
                          client_ip=self.request.getpeername()[0],
                          client_port=self.request.getpeername()[1],
                      ))
              self.request.sendall(resp)
  
  
  class TestHostDBBadResolvConf(helpers.EnvironmentCase):
      '''
      Test that ATS can handle an empty resolv_conf
      '''
      @classmethod
      def setUpEnv(cls, env):
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 2,
              'proxy.config.dns.resolv_conf': '/tmp/non_existant_file',
              'proxy.config.url_remap.remap_required': 0,
  
          })
  
      def test_working(self):
          ret = requests.get('http://trafficserver.readthedocs.org',
                             proxies=self.proxies,
                             )
          self.assertEqual(ret.status_code, 502)
  
  
  class TestHostDBPartiallyFailedDNS(helpers.EnvironmentCase):
      '''
      Tests for how hostdb handles when there is one failed and one working resolver
      '''
      @classmethod
      def setUpEnv(cls, env):
          # TODO: Fix this!
          # This intermittently fails on Jenkins (such as https://ci.trafficserver.apache.org/job/tsqa-master/387/testReport/test_hostdb/TestHostDBPartiallyFailedDNS/test_working/)
          # we aren't sure if this is a failure of ATS or just a race on jenkins (since its slow)
          raise helpers.unittest.SkipTest()
  
          resolv_conf_path = os.path.join(env.layout.prefix, 'resolv.conf')
  
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 2,
              'proxy.config.dns.resolv_conf': resolv_conf_path,
              'proxy.config.url_remap.remap_required': 0,
  
          })
  
          with open(resolv_conf_path, 'w') as fh:
              fh.write('nameserver 1.1.1.0\n')  # some non-existant nameserver
              fh.write('nameserver 8.8.8.8\n')  # some REAL nameserver
  
      def test_working(self):
          start = time.time()
          ret = requests.get('http://trafficserver.readthedocs.org',
                             proxies=self.proxies,
                             )
          self.assertLess(time.time() - start, self.configs['records.config']['CONFIG']['proxy.config.hostdb.lookup_timeout'])
          self.assertEqual(ret.status_code, 200)
  
  
  class TestHostDBFailedDNS(helpers.EnvironmentCase):
      '''
      Tests for how hostdb handles when there is no reachable resolver
      '''
      @classmethod
      def setUpEnv(cls, env):
          resolv_conf_path = os.path.join(env.layout.prefix, 'resolv.conf')
  
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 2,
              'proxy.config.dns.resolv_conf': resolv_conf_path,
              'proxy.config.url_remap.remap_required': 0,
  
          })
  
          with open(resolv_conf_path, 'w') as fh:
              fh.write('nameserver 1.1.1.0\n')  # some non-existant nameserver
  
      def test_lookup_timeout(self):
          start = time.time()
          ret = requests.get('http://some_nonexistant_domain',
                             proxies=self.proxies,
                             )
          self.assertGreater(time.time() - start, self.configs['records.config']['CONFIG']['proxy.config.hostdb.lookup_timeout'])
          self.assertEqual(ret.status_code, 502)
          self.assertIn('ATS', ret.headers['server'])
  
  
  class TestHostDBHostsFile(helpers.EnvironmentCase, tsqa.test_cases.HTTPBinCase):
      '''
      Tests for hostdb's host-file implementation
      '''
      @classmethod
      def setUpEnv(cls, env):
          cls.hosts_file_path = os.path.join(env.layout.prefix, 'hosts')
          with open(cls.hosts_file_path, 'w') as fh:
              fh.write('127.0.0.1 local\n')
              fh.write('127.0.0.2 local2\n')
  
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 2,
              'proxy.config.url_remap.remap_required': 1,
              'proxy.config.http.connect_attempts_max_retries': 1,
              'proxy.config.hostdb.host_file.interval': 1,
              'proxy.config.hostdb.host_file.path': cls.hosts_file_path,
              'proxy.config.diags.debug.enabled': 1,
              'proxy.config.diags.debug.tags': 'hostdb',
          })
          # create a socket server
          cls.socket_server = tsqa.endpoint.SocketServerDaemon(EchoServerIpHandler)
          cls.socket_server.start()
          cls.socket_server.ready.wait()
          cls.configs['remap.config'].add_line('map http://local/ http://local:{0}/'.format(cls.socket_server.port))
          cls.configs['remap.config'].add_line('map http://local2/ http://local2:{0}/'.format(cls.socket_server.port))
          cls.configs['remap.config'].add_line('map http://local3/ http://local3:{0}/'.format(cls.socket_server.port))
  
      def test_basic(self):
          '''
          Test basic fnctionality of hosts files
          '''
          # TODO add stat, then wait for the stat to increment
          time.sleep(2)  # wait for the continuation to load the hosts file
          ret = requests.get(
              'http://local/get',
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 200)
          self.assertEqual('127.0.0.1', ret.headers['X-Server-Ip'])
  
          ret = requests.get(
              'http://local2/get',
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 200)
          self.assertEqual('127.0.0.2', ret.headers['X-Server-Ip'])
  
      def test_reload(self):
          '''
          Test that changes to hosts file get loaded within host_file.interval
          '''
          # TODO add stat, then wait for the stat to increment
          time.sleep(2)  # wait for the continuation to load the hosts file
          ret = requests.get(
              'http://local3/get',
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 502)
  
          with open(self.hosts_file_path, 'a') as fh:
              fh.write('127.0.0.3 local3\n')
  
          # TODO add stat, then wait for the stat to increment, with a timeout
          time.sleep(2)
  
          ret = requests.get(
              'http://local3/get',
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 200)
          self.assertEqual('127.0.0.3', ret.headers['X-Server-Ip'])
  
  
  class TestHostDB(helpers.EnvironmentCase, tsqa.test_cases.HTTPBinCase):
      @classmethod
      def setUpEnv(cls, env):
          cls.dns_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
          cls.dns_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
          cls.dns_sock.bind(('', 0))  # bind to all interfaces on an ephemeral port
          dns_port = cls.dns_sock.getsockname()[1]
  
          # set up dns resolver
          cls.responses = {
              'www.foo.com.': dnslib.server.RR.fromZone("foo.com. 1 A 127.0.0.1"),
              'www.stale_for.com.': dnslib.server.RR.fromZone("foo.com. 1 A 127.0.0.1"),
          }
  
          cls.dns_server = dnslib.server.DNSServer(
              StubDNSResolver(cls.responses),
              port=dns_port,
              address="localhost",
          )
          cls.dns_server.start_thread()
  
          cls.hosts_file_path = os.path.join(env.layout.prefix, 'resolv')
          with open(cls.hosts_file_path, 'w') as fh:
              fh.write('nameserver 127.0.0.1:{0}\n'.format(dns_port))
  
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 1,
              'proxy.config.url_remap.remap_required': 0,
              'proxy.config.http.connect_attempts_max_retries': 1,
              'proxy.config.diags.debug.enabled': 1,
              'proxy.config.diags.debug.tags': 'hostdb',
              'proxy.config.dns.resolv_conf': os.path.join(env.layout.prefix, 'resolv'),
              'proxy.config.hostdb.serve_stale_for': 2,
              'proxy.config.hostdb.ttl_mode': 0,
              'proxy.config.http_ui_enabled': 3,
          })
  
          cls.configs['remap.config'].add_line('map /_hostdb/ http://{hostdb}')
  
      def _hostdb_entries(self):
          # mapping of name -> entries
          ret = {}
          showall_ret = requests.get('http://127.0.0.1:{0}/_hostdb/showall?format=json'.format(
              self.configs['records.config']['CONFIG']['proxy.config.http.server_ports']
          ), timeout=1).json()
  
          for item in showall_ret:
              ret[item['hostname']] = item
  
          return ret
  
      def test_dns(self):
          '''Test that DNS lookups end up in hostdb as we expect
          '''
          # TODO: remove
          self.test_basic()
          print self._hostdb_entries()
  
          # test something with a LARGE number of entries
          zone_parts = []
          # TODO: fix this, right now there is `#define DNS_MAX_ADDRS 35` which
          # controls how many work-- we should make this configurable
          # 30 works, once you pass 35 some records are missing, and at some point
          # you start getting garbage (50 for example) and at some point (100) it
          # seems to crash
          NUM_RECORDS = 2
          for x in xrange(0, NUM_RECORDS):
              zone_parts.append("www.huge.com. 1 A 127.0.0.{0}".format(x + 1))
          self.responses['www.huge.com.'] = dnslib.server.RR.fromZone('\n'.join(zone_parts))
  
          ret = requests.get(
              'http://www.huge.com:{0}/get'.format(self.http_endpoint.address[1]),
              proxies=self.proxies,
          )
          #self.assertEqual(ret.status_code, 200)
  
          for item in self._hostdb_entries()['www.huge.com']['rr_records']:
              print item['ip']
  
          self.assertEqual(len(self._hostdb_entries()['www.huge.com']['rr_records']), NUM_RECORDS)
  
      def test_basic(self):
          '''
          Test basic fnctionality of resolver
          '''
  
          # test one that works
          ret = requests.get(
              'http://www.foo.com:{0}/get'.format(self.http_endpoint.address[1]),
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 200)
  
          # check one that doesn't exist
          ret = requests.get(
              'http://www.bar.com:{0}/get'.format(self.http_endpoint.address[1]),
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 502)
  
      def test_serve_stail_for(self):
          start = time.time()
          ret = requests.get(
              'http://www.stale_for.com:{0}/get'.format(self.http_endpoint.address[1]),
              proxies=self.proxies,
          )
          self.assertEqual(ret.status_code, 200)
          # mark the DNSServer down
          with kill_dns(self.dns_server):
              timeout_at = time.time() + 10
              end_working = None
              end = None
              count = 0
  
              while time.time() < timeout_at:
                  ret = requests.get(
                      'http://www.stale_for.com:{0}/get'.format(self.http_endpoint.address[1]),
                      proxies=self.proxies,
                  )
                  count += 1
                  if ret.status_code != 200:
                      end = time.time()
                      break
                  else:
                      end_working = time.time()
                  time.sleep(0.5)
              # ensure that it was for at least 2 seconds
              print end_working - start
              self.assertTrue(end_working - start >= 2)
              # TODO: Fix this!
              # for whatever reason the failed DNS response is taking ~3.5s to timeout
              # even though the hostdb.lookup_timeout is set to 1 (meaning it should be ~1s)
              # print end - end_working
              #self.assertTrue(end - start >= 2)
  
  
  class TestHostDBSRV(helpers.EnvironmentCase):
      '''Tests for SRV records within hostdb
  
          Tests:
              - SRV record
                  - port overriding
                  - http/https lookups
              - fallback to non SRV
      '''
      SS_CONFIG = {
          '_http._tcp.www.foo.com.': lambda: tsqa.endpoint.SocketServerDaemon(EchoServerIpHandler),
          '_https._tcp.www.foo.com.': lambda: tsqa.endpoint.SSLSocketServerDaemon(
              EchoServerIpHandler,
              helpers.tests_file_path('cert.pem'),
              helpers.tests_file_path('key.pem'),
          ),
      }
  
      # TODO: clean up and push into trafficserver-qa
      @property
      def proxies(self):
          '''
          Return a dict of schema -> proxy. This is primarily used for requests
          '''
          ret = {}
          for item in self.configs['records.config']['CONFIG']['proxy.config.http.server_ports'].split():
              # TODO: better parsing? For now assume first is port
              parts = item.split(':')
              dst = 'http://127.0.0.1:{0}'.format(parts[0])
              if len(parts) > 1:
                  if parts[1] not in ret:
                      ret[parts[1]] = dst
              elif 'http' not in ret:
                  ret['http'] = dst
          return ret
  
      @classmethod
      def setUpEnv(cls, env):
          cls.dns_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
          cls.dns_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
          cls.dns_sock.bind(('', 0))  # bind to all interfaces on an ephemeral port
          dns_port = cls.dns_sock.getsockname()[1]
  
          # set up dns resolver
          cls.responses = {
              'www.foo.com.': dnslib.server.RR.fromZone("foo.com. 1 A 127.0.0.3\nfoo.com. 1 A 127.0.0.2"),
              'www.stale_for.com.': dnslib.server.RR.fromZone("foo.com. 1 A 127.0.0.1"),
          }
  
          cls.dns_server = dnslib.server.DNSServer(
              StubDNSResolver(cls.responses),
              port=dns_port,
              address="localhost",
          )
          cls.dns_server.start_thread()
  
          cls.ssl_port = tsqa.utils.bind_unused_port()[1]
          cls.configs['records.config']['CONFIG']['proxy.config.http.server_ports'] += ' {0}:ssl'.format(cls.ssl_port)
  
          cls.configs['records.config']['CONFIG'].update({
              'proxy.config.http.response_server_enabled': 2,  # only add server headers when there weren't any
              'proxy.config.hostdb.lookup_timeout': 1,
              'proxy.config.http.connect_attempts_max_retries': 1,
              'proxy.config.diags.debug.enabled': 1,
              'proxy.config.diags.debug.tags': 'hostdb',
              'proxy.config.dns.resolv_conf': os.path.join(env.layout.prefix, 'resolv'),
              'proxy.config.hostdb.serve_stale_for': 2,
              'proxy.config.hostdb.ttl_mode': 0,
              'proxy.config.http_ui_enabled': 3,
              'proxy.config.dns.nameservers': '127.0.0.1:{0}'.format(dns_port),
              'proxy.config.srv_enabled': 1,
          })
  
          cls.configs['ssl_multicert.config'].add_line('dest_ip=* ssl_cert_name={0}'.format(
              helpers.tests_file_path('rsa_keys/www.test.com.pem'),
          ))
  
          y = -1
          for name, factory in cls.SS_CONFIG.iteritems():
              y += 1
              ss_dns_results = []
              for x in xrange(0, 3):
                  ss = factory()
                  ss.start()
                  ss.ready.wait()
                  ss_dns_results.append(dnslib.server.RR(
                      name,
                      dnslib.dns.QTYPE.SRV,
                      rdata=dnslib.dns.SRV(
                          priority=10,
                          weight=10,
                          port=ss.port,
                          target='127.0.{0}.{1}.'.format(y, x + 1),  # note: NUM_REALS must be < 253
                      ),
                      ttl=1,
                  ))
              cls.responses[name] = ss_dns_results
  
          cls.configs['remap.config'].add_line('map http://www.foo.com/ http://www.foo.com/')
          cls.configs['remap.config'].add_line('map https://www.foo.com/ https://www.foo.com/')
          cls.configs['remap.config'].add_line('map /_hostdb/ http://{hostdb}')
  
      def _hostdb_entries(self):
          # mapping of name -> entries
          ret = {}
          showall_ret = requests.get('http://127.0.0.1:{0}/_hostdb/showall?format=json'.format(
              self.configs['records.config']['CONFIG']['proxy.config.http.server_ports']
          ), timeout=1)
          return showall_ret.text
  
          for item in showall_ret:
              ret[item['hostname']] = item
  
          return ret
  
      def test_https(self):
          '''Test https SRV lookups
  
          we expect the SRV lookup to get different hosts, but otherwise act the same
          '''
          time.sleep(1)
          expected_set = set([d.rdata.port for d in self.responses['_https._tcp.www.foo.com.']])
  
          actual_set = set()
          for x in xrange(0, 10):
              # test one that works
              ret = requests.get(
                  'https://localhost:{0}/'.format(self.ssl_port),
                  headers={'Host': 'www.foo.com'},
                  verify=False,  # self signed certs, don't bother verifying
              )
              self.assertEqual(ret.status_code, 200)
              actual_set.add(int(ret.headers['X-Server-Port']))
  
          self.assertEqual(expected_set, actual_set)
  
      def test_ports(self):
          '''Test port functionality of SRV responses
  
          SRV responses include ports-- so we want to ensure that we are correctly
          overriding the port based on the response
          '''
          time.sleep(1)
          expected_set = set([d.rdata.port for d in self.responses['_http._tcp.www.foo.com.']])
  
          actual_set = set()
          for x in xrange(0, 10):
              # test one that works
              ret = requests.get(
                  'http://www.foo.com/',
                  proxies=self.proxies,
              )
              self.assertEqual(ret.status_code, 200)
              actual_set.add(int(ret.headers['X-Server-Port']))
  
          self.assertEqual(expected_set, actual_set)
  
      def test_priority(self):
          '''Test port functionality of SRV responses
  
          SRV responses include ports-- so we want to ensure that we are correctly
          overriding the port based on the response
          '''
          time.sleep(3)  # TODO: clear somehow? waiting for expiry is lame
  
          NUM_REQUESTS = 10
          orig_responses = self.responses['_http._tcp.www.foo.com.']
          try:
              self.responses['_http._tcp.www.foo.com.'][0].rdata.priority = 1
  
              request_distribution = {}
              for x in xrange(0, NUM_REQUESTS):
                  # test one that works
                  ret = requests.get(
                      'http://www.foo.com/',
                      proxies=self.proxies,
                  )
                  self.assertEqual(ret.status_code, 200)
                  port = int(ret.headers['X-Server-Port'])
                  if port not in request_distribution:
                      request_distribution[port] = 0
                  request_distribution[port] += 1
  
              # since one has a lower priority, we want to ensure that it got all requests
              self.assertEqual(
                  request_distribution[self.responses['_http._tcp.www.foo.com.'][0].rdata.port],
                  NUM_REQUESTS,
              )
  
          finally:
              self.responses['_http._tcp.www.foo.com.'] = orig_responses
  
      def test_weight(self):
          '''Test port functionality of SRV responses
  
          SRV responses include ports-- so we want to ensure that we are correctly
          overriding the port based on the response
          '''
          time.sleep(3)  # TODO: clear somehow? waiting for expiry is lame
  
          NUM_REQUESTS = 100
          orig_responses = self.responses['_http._tcp.www.foo.com.']
          try:
              self.responses['_http._tcp.www.foo.com.'][0].rdata.weight = 100
  
              request_distribution = {}
              for x in xrange(0, NUM_REQUESTS):
                  # test one that works
                  ret = requests.get(
                      'http://www.foo.com/',
                      proxies=self.proxies,
                  )
                  self.assertEqual(ret.status_code, 200)
                  port = int(ret.headers['X-Server-Port'])
                  if port not in request_distribution:
                      request_distribution[port] = 0
                  request_distribution[port] += 1
  
              # since the first one has a significantly higher weight, we expect it to
              # take ~10x the traffic of the other 2
              self.assertTrue(
                  request_distribution[self.responses['_http._tcp.www.foo.com.'][0].rdata.port] >
                  (NUM_REQUESTS / len(self.responses['_http._tcp.www.foo.com.'])) * 2,
                  'Expected significantly more traffic on {0} than the rest: {1}'.format(
                      self.responses['_http._tcp.www.foo.com.'][0].rdata.port,
                      request_distribution,
                  ),
              )
  
          finally:
              self.responses['_http._tcp.www.foo.com.'] = orig_responses
  