text: # Copyright: 2019, NLnet Labs and the Internet.nl contributors
  # SPDX-License-Identifier: Apache-2.0
  from collections import OrderedDict
  import socket
  import time
  
  import pythonwhois
  import sys
  import unbound
  
  from celery import shared_task
  from celery.exceptions import SoftTimeLimitExceeded
  from django.conf import settings
  from django.core.cache import cache
  from django.db import transaction
  
  from . import SetupUnboundContext
  from . import shared
  from .dispatcher import check_registry, post_callback_hook
  from .. import batch, batch_shared_task
  from .. import scoring, categories, redis_id
  from ..models import DomainTestDnssec, DnssecStatus, MailTestDnssec
  
  
  UNBOUND_PATCHED_DS_LOG = "internetnl - DS unsupported"
  
  
  @shared_task(bind=True)
  def web_callback(self, results, addr, req_limit_id):
      category = categories.WebDnssec()
      dtdnssec = save_results_web(addr, results, category)
      # Always calculate scores on saving.
      from ..probes import web_probe_dnssec
      web_probe_dnssec.rated_results_by_model(dtdnssec)
      post_callback_hook(req_limit_id, self.request.id)
      return results
  
  
  @batch_shared_task(bind=True)
  def batch_web_callback(self, results, addr):
      category = categories.WebDnssec()
      dtdnssec = save_results_web(addr, results, category)
      # Always calculate scores on saving.
      from ..probes import batch_web_probe_dnssec
      batch_web_probe_dnssec.rated_results_by_model(dtdnssec)
      batch.scheduler.batch_callback_hook(dtdnssec, self.request.id)
  
  
  @shared_task(bind=True)
  def mail_callback(self, results, addr, req_limit_id):
      category = categories.MailDnssec()
      maildomain = save_results_mail(addr, results, category)
      # Always calculate scores on saving.
      from ..probes import mail_probe_dnssec
      mail_probe_dnssec.rated_results_by_model(maildomain)
      post_callback_hook(req_limit_id, self.request.id)
      return results
  
  
  @batch_shared_task(bind=True)
  def batch_mail_callback(self, results, addr):
      category = categories.MailDnssec()
      maildomain = save_results_mail(addr, results, category)
      # Always calculate scores on saving.
      from ..probes import batch_mail_probe_dnssec
      batch_mail_probe_dnssec.rated_results_by_model(maildomain)
      batch.scheduler.batch_callback_hook(maildomain, self.request.id)
  
  
  web_registered = check_registry("dnssec", web_callback)
  mail_registered = check_registry(
      "mail_dnssec", mail_callback, shared.mail_get_servers)
  batch_web_registered = check_registry("batch_dnssec", batch_web_callback)
  batch_mail_registered = check_registry(
      "batch_mail_dnssec", batch_mail_callback, shared.batch_mail_get_servers)
  
  
  @web_registered
  @shared_task(
      bind=True,
      soft_time_limit=settings.SHARED_TASK_SOFT_TIME_LIMIT_LOW,
      time_limit=settings.SHARED_TASK_TIME_LIMIT_LOW,
      base=SetupUnboundContext)
  def web_is_secure(self, url, *args, **kwargs):
      return do_web_is_secure(self, url, *args, **kwargs)
  
  
  @batch_web_registered
  @batch_shared_task(
      bind=True,
      soft_time_limit=settings.BATCH_SHARED_TASK_SOFT_TIME_LIMIT_HIGH,
      time_limit=settings.BATCH_SHARED_TASK_TIME_LIMIT_HIGH,
      base=SetupUnboundContext)
  def batch_web_is_secure(self, url, *args, **kwargs):
      return do_web_is_secure(self, url, *args, **kwargs)
  
  
  @mail_registered
  @shared_task(
      bind=True,
      soft_time_limit=settings.SHARED_TASK_SOFT_TIME_LIMIT_LOW,
      time_limit=settings.SHARED_TASK_TIME_LIMIT_LOW,
      base=SetupUnboundContext)
  def mail_is_secure(self, mailservers, url, *args, **kwargs):
      return do_mail_is_secure(self, mailservers, url, *args, **kwargs)
  
  
  @batch_mail_registered
  @batch_shared_task(
      bind=True,
      soft_time_limit=settings.BATCH_SHARED_TASK_SOFT_TIME_LIMIT_HIGH,
      time_limit=settings.BATCH_SHARED_TASK_TIME_LIMIT_HIGH,
      base=SetupUnboundContext)
  def batch_mail_is_secure(self, mailservers, url, *args, **kwargs):
      return do_mail_is_secure(self, mailservers, url, *args, **kwargs)
  
  
  def registrar_lookup(addr):
      """
      Get the registrar from WHOIS.
  
      Cache the result to avoid being ratelimited too quickly.
  
      """
      res = ""
      if (
          'pythonwhois' in sys.modules
          and not settings.ENABLE_BATCH
      ):
          cache_id = redis_id.whois.id.format(addr)
          cache_ttl = redis_id.whois.ttl
          cached = cache.get(cache_id)
          if cached:
              res = cached
          else:
              try:
                  whois = pythonwhois.get_whois(".".join(addr.split(".")[-2:]))
                  if (whois and isinstance(whois, dict)
                          and whois.get("registrar")):
                      res = ", ".join(whois["registrar"])[:250]
              except (
                      socket.error, pythonwhois.shared.WhoisException,
                      UnicodeDecodeError, IndexError):
                  pass
  
              cache.set(cache_id, res, cache_ttl)
      return res
  
  
  @transaction.atomic
  def save_results_mail(addr, results, category):
      """
      Save results in the DB for the mail test.
  
      """
      mailtdnssec = MailTestDnssec(domain=addr)
      mailtdnssec.save()
      testname, result = results[0]
      domain_report = {}
      subreports = {}
      # For the mail test the first result is the domain, the rest are the MX.
      for i, (domain, r) in enumerate(result.items()):
          category = category.__class__()
          dtdnssec = DomainTestDnssec(domain=domain)
          if i == 0:
              report, status, score, log = get_domain_results(
                  domain, r, category)
              # We get the domain results here, domain does not belong to a
              # summary report like the MX.
              domain_report['dnssec_exists'] = report['dnssec_exists']
              domain_report['dnssec_valid'] = report['dnssec_valid']
          else:
              report, status, score, log = get_mx_results(r, category)
              subreports[domain] = report
          dtdnssec.report = report
          dtdnssec.status = status
          dtdnssec.score = score
          dtdnssec.log = log
          dtdnssec.save()
          mailtdnssec.testset.add(dtdnssec)
  
      # Build the summary report for the MX.
      category = category.__class__()
      if i == 0:
          # Give a warning for the mailserver part of the DNSSEC test when there
          # are no mailservers.
          category.subtests['dnssec_mx_exists'].result_no_mailservers()
      mx_report = category.gen_report()
      shared.aggregate_subreports(subreports, mx_report)
  
      mx_report.update(domain_report)
      mailtdnssec.report = mx_report
      mailtdnssec.save()
      return mailtdnssec
  
  
  def save_results_web(addr, results, category):
      """
      Save results in the DB for the web test.
  
      """
      dtdnssec = DomainTestDnssec(domain=addr)
      testname, result = results[0]
      # For the web test there is going to be only one item in the result dict.
      for domain, r in result.items():
          report, status, score, log = get_domain_results(domain, r, category)
          dtdnssec.report = report
          dtdnssec.status = status
          dtdnssec.score = score
          dtdnssec.log = log
          dtdnssec.save()
      return dtdnssec
  
  
  def get_domain_results(domain, r, category):
      status = r.get('status')
      score = r.get('score')
      log = ""
      registrar = registrar_lookup(domain)
      if status == DnssecStatus.secure.value:
          status = DnssecStatus.secure
          category.subtests['dnssec_exists'].result_good(
              [[domain, registrar]])
  
          category.subtests['dnssec_valid'].result_good(
              [[domain, "detail tech data secure"]])
  
      elif status == DnssecStatus.insecure.value:
          status = DnssecStatus.insecure
          log = r.get('log')
          unsupported_ds_algo = UNBOUND_PATCHED_DS_LOG in log
          if unsupported_ds_algo:
              category.subtests['dnssec_exists'].result_good(
                  [[domain, registrar]])
  
              category.subtests['dnssec_valid'].result_unsupported_ds_algo(
                  [[domain, "detail tech data insecure"]])
          else:
              log = ""  # Don't store the log for simple insecure.
              category.subtests['dnssec_exists'].result_bad(
                  [[domain, registrar]])
  
              category.subtests['dnssec_valid'].result_insecure(
                  [[domain, "detail tech data insecure"]])
  
      elif status == DnssecStatus.bogus.value:
          status = DnssecStatus.bogus
          log = r.get('log')
          category.subtests['dnssec_exists'].result_good(
              [[domain, registrar]])
  
          category.subtests['dnssec_valid'].result_bad(
              [[domain, "detail tech data bogus"]])
  
      elif status == DnssecStatus.servfail.value:
          status = DnssecStatus.servfail
          category.subtests['dnssec_exists'].result_servfail(
              [[domain, registrar]])
  
          category.subtests['dnssec_valid'].result_servfail(
              [[domain, "detail tech data not-tested"]])
  
      else:
          status = DnssecStatus.dnserror
          category.subtests['dnssec_exists'].result_resolver_error(
              [[domain, registrar]])
  
          category.subtests['dnssec_valid'].result_resolver_error(
              [[domain, "detail tech data not-tested"]])
  
      report = category.gen_report()
  
      return report, status, score, log
  
  
  def get_mx_results(results, category):
      status = results.get("status")
      score = results.get("score")
      log = ""
      if status == DnssecStatus.secure.value:
          status = DnssecStatus.secure
          category.subtests['dnssec_mx_exists'].result_good()
          category.subtests['dnssec_mx_valid'].result_good()
  
      elif status == DnssecStatus.insecure.value:
          status = DnssecStatus.insecure
          log = results.get('log')
          unsupported_ds_algo = UNBOUND_PATCHED_DS_LOG in log
          if unsupported_ds_algo:
              category.subtests['dnssec_mx_exists'].result_good()
              category.subtests['dnssec_mx_valid'].result_unsupported_ds_algo()
          else:
              log = ""  # Don't store the log for simple insecure.
              category.subtests['dnssec_mx_exists'].result_bad()
              category.subtests['dnssec_mx_valid'].result_insecure()
  
      elif status == DnssecStatus.bogus.value:
          status = DnssecStatus.bogus
          log = results.get('log')
          category.subtests['dnssec_mx_exists'].result_good()
          category.subtests['dnssec_mx_valid'].result_bad()
  
      elif status == DnssecStatus.servfail.value:
          status = DnssecStatus.servfail
          category.subtests['dnssec_mx_exists'].result_servfail()
  
      else:
          status = DnssecStatus.dnserror
          category.subtests['dnssec_mx_exists'].result_resolver_error()
  
      report = category.gen_report()
  
      return report, status, score, log
  
  
  def unbound_callback(data, status, r):
      """
      Check unbound's response and record the status, score and log (if any) in
      the callback data.
  
      """
      data['score'] = data['score_error']
      data['log'] = ""
      if status == 0:
          log = ""
          if r.rcode == unbound.RCODE_SERVFAIL:
              status = DnssecStatus.servfail.value
              score = data['score_error']
          elif r.secure:
              status = DnssecStatus.secure.value
              score = data['score_secure']
          elif r.bogus:
              status = DnssecStatus.bogus.value
              log = r.why_bogus
              score = data['score_bogus']
          else:
              status = DnssecStatus.insecure.value
              log = r.why_bogus
              score = data['score_insecure']
          data['score'] = score
          data['log'] = log or ""
          data['status'] = status
      data['done'] = True
  
  
  def do_web_is_secure(self, url, *args, **kwargs):
      try:
          dnssec_result = dnssec_status(
              url, self,
              score_secure=scoring.WEB_DNSSEC_SECURE,
              score_insecure=scoring.WEB_DNSSEC_INSECURE,
              score_bogus=scoring.WEB_DNSSEC_BOGUS,
              score_error=scoring.WEB_DNSSEC_ERROR)
  
      except SoftTimeLimitExceeded:
          dnssec_result = dict(
              status=DnssecStatus.dnserror.value,
              score=scoring.WEB_DNSSEC_ERROR,
              log="Timed out")
  
      return ('is_secure', {url: dnssec_result})
  
  
  def do_mail_is_secure(self, mailservers, url, *args, **kwargs):
      try:
          mailservers.insert(0, (url, None))
          res = OrderedDict()
          for domain, _ in mailservers:
              if domain != "":
                  res[domain] = dnssec_status(
                      domain, self,
                      score_secure=scoring.MAIL_DNSSEC_SECURE,
                      score_insecure=scoring.MAIL_DNSSEC_INSECURE,
                      score_bogus=scoring.MAIL_DNSSEC_BOGUS,
                      score_error=scoring.MAIL_DNSSEC_ERROR)
  
      except SoftTimeLimitExceeded:
          for domain, _ in mailservers:
              if domain != "" and not res.get(domain):
                  res[domain] = dict(
                      status=DnssecStatus.dnserror.value,
                      score=scoring.MAIL_DNSSEC_ERROR,
                      log="Timed out")
  
      return ('is_secure', res)
  
  
  def dnssec_status(
          domain, ub_task, score_secure, score_insecure, score_bogus,
          score_error):
      """
      Check the DNSSEC status of the domain.
  
      """
      try:
          cb_data = dict(
              done=False, status=None, log="", score=score_error,
              score_secure=score_secure, score_insecure=score_insecure,
              score_bogus=score_bogus, score_error=score_error)
          retval, async_id = ub_task.ub_ctx.resolve_async(
              domain, cb_data, unbound_callback, unbound.RR_TYPE_SOA,
              unbound.RR_CLASS_IN)
          status = DnssecStatus.dnserror.value
          if retval == 0:
              while not cb_data["done"]:
                  time.sleep(0.1)
                  retval = ub_task.ub_ctx.process()
                  if retval != 0:
                      cb_data["status"] = DnssecStatus.dnserror.value
                      break
              status = cb_data["status"]
  
      except SoftTimeLimitExceeded as e:
          if async_id:
              ub_task.ub_ctx.cancel(async_id)
          raise e
  
      return dict(
          status=status,
          score=cb_data["score"],
          log=cb_data["log"])
  