From 8f339923f87a989c58589251856e207a73e684d1 Mon Sep 17 00:00:00 2001
From: Carlos Quintana <carlos@cquintana.dev>
Date: Thu, 24 Feb 2022 15:05:05 +0100
Subject: [PATCH] Make nameservers configurable

---
 app/config.py    | 17 ++++++++++++++++-
 app/dns_utils.py | 15 +++++++--------
 example.env      |  6 +++++-
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/app/config.py b/app/config.py
index 2148eb42..7578f01f 100644
--- a/app/config.py
+++ b/app/config.py
@@ -4,7 +4,7 @@ import socket
 import string
 import subprocess
 from ast import literal_eval
-from typing import Callable, List
+from typing import Callable, List, Optional
 from urllib.parse import urlparse
 
 from dotenv import load_dotenv
@@ -37,6 +37,13 @@ def sl_getenv(env_var: str, default_factory: Callable = None):
     return literal_eval(value)
 
 
+def env(var: str, default: Optional[str] = None) -> Optional[str]:
+    if var in os.environ:
+        return os.environ[var]
+    else:
+        return default
+
+
 config_file = os.environ.get("CONFIG")
 if config_file:
     config_file = get_abs_path(config_file)
@@ -431,3 +438,11 @@ def get_allowed_redirect_domains() -> List[str]:
 
 
 ALLOWED_REDIRECT_DOMAINS = get_allowed_redirect_domains()
+
+
+def setup_nameservers():
+    nameservers = env("NAMESERVERS", "1.1.1.1")
+    return nameservers.split(",")
+
+
+NAMESERVERS = setup_nameservers()
diff --git a/app/dns_utils.py b/app/dns_utils.py
index 0c27b939..80e25043 100644
--- a/app/dns_utils.py
+++ b/app/dns_utils.py
@@ -1,3 +1,4 @@
+from app import config
 from typing import Optional, List, Tuple
 
 import dns.resolver
@@ -5,16 +6,14 @@ import dns.resolver
 
 def _get_dns_resolver():
     my_resolver = dns.resolver.Resolver()
-
-    # 1.1.1.1 is CloudFlare's public DNS server
-    my_resolver.nameservers = ["1.1.1.1"]
+    my_resolver.nameservers = config.NAMESERVERS
 
     return my_resolver
 
 
 def get_ns(hostname) -> [str]:
     try:
-        answers = _get_dns_resolver().resolve(hostname, "NS")
+        answers = _get_dns_resolver().resolve(hostname, "NS", search=True)
     except Exception:
         return []
     return [a.to_text() for a in answers]
@@ -23,7 +22,7 @@ def get_ns(hostname) -> [str]:
 def get_cname_record(hostname) -> Optional[str]:
     """Return the CNAME record if exists for a domain, WITHOUT the trailing period at the end"""
     try:
-        answers = _get_dns_resolver().resolve(hostname, "CNAME")
+        answers = _get_dns_resolver().resolve(hostname, "CNAME", search=True)
     except Exception:
         return None
 
@@ -39,7 +38,7 @@ def get_mx_domains(hostname) -> [(int, str)]:
     domain name ends with a "." at the end.
     """
     try:
-        answers = _get_dns_resolver().resolve(hostname, "MX")
+        answers = _get_dns_resolver().resolve(hostname, "MX", search=True)
     except Exception:
         return []
 
@@ -60,7 +59,7 @@ _include_spf = "include:"
 def get_spf_domain(hostname) -> [str]:
     """return all domains listed in *include:*"""
     try:
-        answers = _get_dns_resolver().resolve(hostname, "TXT")
+        answers = _get_dns_resolver().resolve(hostname, "TXT", search=True)
     except Exception:
         return []
 
@@ -82,7 +81,7 @@ def get_spf_domain(hostname) -> [str]:
 def get_txt_record(hostname) -> [str]:
     """return all domains listed in *include:*"""
     try:
-        answers = _get_dns_resolver().resolve(hostname, "TXT")
+        answers = _get_dns_resolver().resolve(hostname, "TXT", search=True)
     except Exception:
         return []
 
diff --git a/example.env b/example.env
index e3b9cadb..31429f75 100644
--- a/example.env
+++ b/example.env
@@ -174,4 +174,8 @@ DISABLE_ONBOARDING=true
 #ALIAS_AUTOMATIC_DISABLE=true
 
 # domains that can be present in the &next= section when using absolute urls
-ALLOWED_REDIRECT_DOMAINS=[]
\ No newline at end of file
+ALLOWED_REDIRECT_DOMAINS=[]
+
+# DNS nameservers to be used by the app
+# Multiple nameservers can be specified, separated by ','
+NAMESERVERS="1.1.1.1"